Update app.py

app.py

@@ -3,6 +3,7 @@ from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
 import spaces
 import torch
 import re
+from PIL import Image
 
 model = PaliGemmaForConditionalGeneration.from_pretrained("gokaygokay/sd3-long-captioner").to("cpu").eval()
 processor = PaliGemmaProcessor.from_pretrained("gokaygokay/sd3-long-captioner")
@@ -28,23 +29,31 @@ def modify_caption(caption: str) -> str:
 
     return re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE)
 
-def create_captions_rich(images):
+def create_captions_rich(files):
     captions = []
     prompt = "caption en"
     
-    for image in images:
+    for file_path in files:
+        try:
+            image = Image.open(file_path.name)
+        except Exception as e:
+            captions.append(f"Error opening image: {e}")
+            continue
+        
         model_inputs = processor(text=prompt, images=image, return_tensors="pt").to("cpu")
         input_len = model_inputs["input_ids"].shape[-1]
         
-        with torch.inference_mode():
-            generation = model.generate(**model_inputs, max_new_tokens=256, do_sample=False)
-            generation = generation[0][input_len:]
-            decoded = processor.decode(generation, skip_special_tokens=True)
-
-            modified_caption = modify_caption(decoded)
-            captions.append(modified_caption)
+        try:
+            with torch.no_grad():
+                generation = model.generate(**model_inputs, max_new_tokens=256, do_sample=False)
+                generation = generation[0][input_len:]
+                decoded = processor.decode(generation, skip_special_tokens=True)
+                modified_caption = modify_caption(decoded)
+                captions.append(modified_caption)
+        except Exception as e:
+            captions.append(f"Error generating caption: {e}")
     
-    return captions
+    return "\n".join(captions)
 
 css = """
   #mkd {
@@ -60,10 +69,10 @@ with gr.Blocks(css=css) as demo:
     with gr.Tab(label="Image to Prompt for SD3"):
         with gr.Row():
             with gr.Column():
-                input_imgs = gr.Image(label="Input Images", type="pil", interactive=True)
+                input_files = gr.Files(label="Input Images")
                 submit_btn = gr.Button(value="Start")
             outputs = gr.Textbox(label="Prompts", lines=10, interactive=False)
 
-        submit_btn.click(create_captions_rich, inputs=[input_imgs], outputs=[outputs])
+        submit_btn.click(create_captions_rich, inputs=[input_files], outputs=[outputs])
 
 demo.launch(debug=True)