Update app.py

app.py

@@ -2,7 +2,7 @@
 import os, io, re, sys, subprocess, hashlib, pathlib, time
 from typing import Optional
 import requests
-from PIL import Image, ImageSequence
+from PIL import Image, ImageSequence, UnidentifiedImageError # Import UnidentifiedImageError
 import gradio as gr
 
 MODEL_DIR = pathlib.Path("model")
@@ -52,11 +52,24 @@ def mp4_to_gif(mp4_bytes: bytes) -> bytes:
     r2 = requests.get(gif_url, timeout=60); r2.raise_for_status(); return r2.content
 
 def load_first_frame(raw: bytes):
-    img = Image.open(io.BytesIO(raw))
-    if getattr(img, "is_animated", False):
-        img = next(ImageSequence.Iterator(img))
-    if img.mode != "RGB": img = img.convert("RGB")
-    return img
+    # Added specific handling for PIL errors
+    try:
+        img = Image.open(io.BytesIO(raw))
+        if getattr(img, "is_animated", False):
+            # Also wrap ImageSequence.Iterator in case of corrupted animated images
+            try:
+                img = next(ImageSequence.Iterator(img))
+            except Exception as e:
+                raise ValueError(f"Could not extract first frame from animated image: {e}")
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+        return img
+    except UnidentifiedImageError:
+        raise ValueError("Could not identify image format or image is corrupted.")
+    except Exception as e:
+        # Catch other PIL errors (e.g., IOErrors during parsing)
+        raise ValueError(f"Failed to load or process image with PIL: {e}")
+
 
 def rebuild_llama_cpp():
     env = os.environ.copy()
@@ -120,19 +133,28 @@ def generate_caption_from_url(url: str, prompt: str="Describe the image."):
                 raw = mp4_to_gif(raw)
             except Exception as e:
                 return "MP4→GIF conversion failed: " + str(e)
-        img = load_first_frame(raw)
-    except Exception as e:
+        img = load_first_frame(raw) # This function now has specific error handling
+    except ValueError as e: # Catch the specific ValueError raised by load_first_frame
         return "Image processing error: " + str(e)
+    except Exception as e: # General fallback for other unexpected image issues
+        return "An unexpected image processing error occurred: " + str(e)
 
+    # Added try-except for image resizing as well
     try:
         img = img.resize((512,512), resample=Image.BICUBIC)
-    except Exception:
-        pass
+    except Exception as e:
+        print(f"Warning: Image resizing failed: {e}. Attempting to proceed without resizing.")
+        # Optionally, you might want to return an error here if resizing is critical.
+        # For captioning, not resizing might just lead to a slightly different result.
 
     # create a tiny base64 tag to signal image presence (model must understand this format)
     import base64
     buf = io.BytesIO()
-    img.save(buf, format="PNG")
+    try:
+        img.save(buf, format="PNG") # Wrap image save in try-except
+    except Exception as e:
+        return "Failed to encode image to base64: " + str(e)
+
     b64 = base64.b64encode(buf.getvalue()).decode()
     img_tag = b64  # minimal
 
@@ -155,4 +177,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch(server_name="0.0.0.0", server_port=7860)