Update app.py

app.py

@@ -5,7 +5,8 @@ import torch
 import random
 import os
 import tempfile
-from PIL import Image
+from PIL import Image, ImageOps
+import pillow_heif  # For HEIF/AVIF support
 
 # Import the pipeline from diffusers
 from diffusers import FluxKontextPipeline
@@ -20,6 +21,9 @@ def load_model():
     """Load the model on CPU first, then move to GPU when needed"""
     global pipe
     if pipe is None:
+        # Register HEIF opener with PIL for AVIF/HEIF support
+        pillow_heif.register_heif_opener()
+        
         # Get token from environment variable
         hf_token = os.getenv("HF_TOKEN")
         if hf_token:
@@ -56,7 +60,17 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
     input_image = None
     if files:
         print(f"Received image: {files[0]}")
-        input_image = Image.open(files[0]).convert("RGB")
+        try:
+            # Try to open and convert the image
+            input_image = Image.open(files[0])
+            # Convert to RGB if needed (handles RGBA, P, etc.)
+            if input_image.mode != "RGB":
+                input_image = input_image.convert("RGB")
+            # Auto-orient the image based on EXIF data
+            input_image = ImageOps.exif_transpose(input_image)
+        except Exception as e:
+            raise gr.Error(f"Could not process the uploaded image: {str(e)}. Please try uploading a different image format (JPEG, PNG, WebP).")
+            
         image = pipe(
             image=input_image,
             prompt=prompt,