Update app.py

app.py

@@ -1,6 +1,7 @@
 import spaces
 import gradio as gr
 import torch
+from PIL import Image
 from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor, pipeline
 from transformers import AutoProcessor, AutoModelForCausalLM
 import re
@@ -56,6 +57,10 @@ MAX_SEED = 2**32 - 1
 
 # Florence caption function
 def florence_caption(image):
+    # Convert image to PIL if it's not already
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    
     inputs = florence_processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
     generated_ids = florence_model.generate(
         input_ids=inputs["input_ids"],
@@ -137,10 +142,13 @@ def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height,
     
     return image, seed
 
-# Gradio Interface
 @spaces.GPU
 def process_workflow(image, text_prompt, vlm_model_choice, use_enhancer, model_choice, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
     if image is not None:
+        # Convert image to PIL if it's not already
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+        
         if vlm_model_choice == "Long Captioner":
             prompt = create_captions_rich(image)
         else:  # Florence