Update app.py

app.py

@@ -1,27 +1,40 @@
 import os
 import io
 import zipfile
-import cv2
-import numpy as np
-from PIL import Image
-from rembg import remove
-from huggingface_hub import InferenceClient
-import gradio as gr
+import sys
+import traceback
+
+# CRITICAL FIX: Catch OpenCV import errors before they crash the app
+try:
+    import cv2
+    import numpy as np
+    from PIL import Image
+    from rembg import remove
+    from huggingface_hub import InferenceClient
+    import gradio as gr
+except Exception as e:
+    print("CRITICAL IMPORT ERROR:", e)
+    print(traceback.format_exc())
+    sys.exit(1)
 
 # Get your HF Token from environment variables
 HF_TOKEN = os.getenv("HF_TOKEN")
-client = InferenceClient(token=HF_TOKEN)
+client = InferenceClient(token=HF_TOKEN, model="stabilityai/stable-diffusion-2-1")
 
 def enhance_lighting(pil_img):
     """Fix exposure and contrast using OpenCV CLAHE"""
-    cv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
-    lab = cv2.cvtColor(cv_img, cv2.COLOR_BGR2LAB)
-    l, a, b = cv2.split(lab)
-    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
-    l = clahe.apply(l)
-    lab = cv2.merge([l, a, b])
-    cv_img = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
-    return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
+    try:
+        cv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
+        lab = cv2.cvtColor(cv_img, cv2.COLOR_BGR2LAB)
+        l, a, b = cv2.split(lab)
+        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
+        l = clahe.apply(l)
+        lab = cv2.merge([l, a, b])
+        cv_img = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
+        return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
+    except Exception as e:
+        print(f"Lighting fix failed: {e}")
+        return pil_img
 
 def fix_aspect_ratio(pil_img, ratio_str):
     """Center-crop image to selected aspect ratio"""
@@ -42,36 +55,41 @@ def process_batch(files, ratio, rm_bg, prompt):
     processed = []
     
     for file in files:
-        img = Image.open(file).convert("RGB")
-        
-        # 1. Fix Lighting
-        img = enhance_lighting(img)
-        
-        # 2. Fix Aspect Ratio
-        img = fix_aspect_ratio(img, ratio)
-        
-        # 3. Remove Background (local rembg on HF Space)
-        if rm_bg:
-            img = remove(img) # Returns RGBA with transparent bg
-        
-        # 4. Generate AI Background (only if both prompt AND bg removal are selected)
-        if prompt and rm_bg:
-            # Generate background image via HF free inference API
-            bg_img = client.text_to_image(
-                f"Product photography background of {prompt}, elegant, soft studio light, photorealistic, 8k",
-                model="stabilityai/stable-diffusion-2-1"
-            )
-            # Resize background to match the jewelery size
-            bg_img = bg_img.resize(img.size).convert("RGBA")
+        try:
+            img = Image.open(file).convert("RGB")
+            
+            # 1. Fix Lighting
+            img = enhance_lighting(img)
             
-            # Compose: paste jewelry over background
-            final = bg_img.copy()
-            final.paste(img, (0, 0), img) # img has alpha channel
-            img = final
+            # 2. Fix Aspect Ratio
+            img = fix_aspect_ratio(img, ratio)
             
-        processed.append(img)
+            # 3. Remove Background (local rembg)
+            if rm_bg:
+                img = remove(img)
+            
+            # 4. Generate AI Background
+            if prompt and rm_bg:
+                print(f"Generating background for: {prompt}")
+                bg_img = client.text_to_image(
+                    f"Product photography background of {prompt}, elegant, soft studio light, photorealistic, 8k",
+                    model="stabilityai/stable-diffusion-2-1"
+                )
+                bg_img = bg_img.resize(img.size).convert("RGBA")
+                
+                final = bg_img.copy()
+                final.paste(img, (0, 0), img)
+                img = final
+                
+            processed.append(img)
+        except Exception as e:
+            print(f"Error processing file {file.name}: {e}")
+            continue
+    
+    if not processed:
+        return None
     
-    # Save all images into a single ZIP for batch download
+    # Save all images into a single ZIP
     zip_buffer = io.BytesIO()
     with zipfile.ZipFile(zip_buffer, "w") as zf:
         for i, img in enumerate(processed):
@@ -83,7 +101,7 @@ def process_batch(files, ratio, rm_bg, prompt):
 
 # Gradio UI Interface
 with gr.Blocks(title="Jewelry Batch Processor") as demo:
-    gr.Markdown("## ✨ Free Jewelry Batch Processor (Powered by Hugging Face)")
+    gr.Markdown("## ✨ Free Jewelry Batch Processor")
     
     with gr.Row():
         files = gr.Files(label="Upload Jewelry Images", file_count="multiple")
@@ -98,4 +116,6 @@ with gr.Blocks(title="Jewelry Batch Processor") as demo:
     
     btn.click(process_batch, inputs=[files, ratio, rm_bg, prompt], outputs=output)
 
-demo.launch()
+# CRITICAL FIX: This ensures Gradio launches correctly on Spaces
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)