Update app.py

app.py

@@ -3,184 +3,152 @@ import torch
 from diffusers import StableDiffusionInpaintPipeline
 from PIL import Image, ImageDraw, ImageFilter
 import numpy as np
-import os
+import spaces
 
-# Clear cache if needed
-os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = '1'
+# Initialize model globally (loaded on CPU first)
+print("Loading model on CPU first (ZeroGPU will move it to GPU when needed)...")
 
-# Device setup
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if device == "cuda" else torch.float32
-print(f"Using device: {device}")
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-inpainting",
+    torch_dtype=torch.float16,
+    safety_checker=None,
+    requires_safety_checker=False
+)
 
-# Try loading model with better error handling
-MODEL_LOADED = False
-pipe = None
+# Don't move to GPU yet - ZeroGPU will handle this
+pipe.enable_attention_slicing()
 
-try:
-    print("Downloading and loading model... This may take a few minutes on first run.")
-    
-    # Use the correct model ID and revision
-    pipe = StableDiffusionInpaintPipeline.from_pretrained(
-        "runwayml/stable-diffusion-inpainting",
-        revision="fp16" if device == "cuda" else "main",
-        torch_dtype=dtype,
-        safety_checker=None,
-        requires_safety_checker=False,
-        use_safetensors=True,  # Force safetensors format
-        local_files_only=False  # Allow downloading
-    )
-    
-    pipe = pipe.to(device)
-    pipe.enable_attention_slicing()
-    
-    MODEL_LOADED = True
-    print("✅ Model loaded successfully!")
-    
-except Exception as e:
-    print(f"❌ First attempt failed: {e}")
-    print("Trying alternative model...")
-    
-    try:
-        # Alternative: Try SD 2.0 inpainting
-        pipe = StableDiffusionInpaintPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-2-inpainting",
-            torch_dtype=dtype,
-            safety_checker=None,
-            requires_safety_checker=False
-        )
-        
-        pipe = pipe.to(device)
-        pipe.enable_attention_slicing()
-        
-        MODEL_LOADED = True
-        print("✅ Alternative model loaded successfully!")
-        
-    except Exception as e2:
-        print(f"❌ Both models failed to load: {e2}")
-        MODEL_LOADED = False
+print("✅ Model loaded! ZeroGPU will activate when generating.")
 
 # Clothing prompts
-CLOTHES = {
-    "Indian Sari": "woman wearing elegant red and gold sari, traditional Indian dress, beautiful saree",
-    "Japanese Kimono": "person wearing beautiful floral kimono, traditional Japanese clothing, silk kimono",
-    "African Dashiki": "person wearing colorful African dashiki with patterns, traditional clothing",
-    "Chinese Qipao": "woman wearing elegant qipao cheongsam dress, traditional Chinese dress",
-    "Scottish Kilt": "man wearing traditional Scottish kilt with tartan pattern",
-    "Middle Eastern Thobe": "person wearing white thobe robe, traditional Middle Eastern clothing"
+CLOTHING_PROMPTS = {
+    "Indian Sari": "woman wearing beautiful red and gold silk sari, traditional Indian saree dress, intricate embroidery",
+    "Japanese Kimono": "person wearing elegant floral kimono with obi belt, traditional Japanese clothing, silk fabric",
+    "African Dashiki": "person wearing colorful African dashiki with geometric patterns, traditional clothing, vibrant",
+    "Chinese Qipao": "woman wearing elegant red qipao cheongsam dress, traditional Chinese dress, silk with gold patterns",
+    "Scottish Kilt": "man wearing traditional Scottish kilt with tartan pattern, highland dress, sporran",
+    "Middle Eastern Thobe": "person wearing white thobe robe, traditional Middle Eastern clothing, flowing fabric"
 }
 
-def create_mask(image):
-    """Create body mask"""
-    w, h = image.size
-    mask = Image.new('L', (w, h), 0)
+def create_body_mask(image):
+    """Create mask for clothing area"""
+    width, height = image.size
+    mask = Image.new('L', (width, height), 0)
     draw = ImageDraw.Draw(mask)
     
     # Body area ellipse
-    left = w * 0.2
-    top = h * 0.25
-    right = w * 0.8
-    bottom = h * 0.9
+    left = width * 0.2
+    top = height * 0.25
+    right = width * 0.8
+    bottom = height * 0.95
     
     draw.ellipse([left, top, right, bottom], fill=255)
     mask = mask.filter(ImageFilter.GaussianBlur(radius=20))
     
     return mask
 
-def process_image(image, clothing_type, steps=20):
-    """Generate traditional clothing"""
+@spaces.GPU(duration=60)  # Request GPU for 60 seconds
+def generate_clothing(input_image, clothing_type, num_steps=25, guidance_scale=7.5):
+    """Generate traditional clothing with ZeroGPU"""
     
-    if not MODEL_LOADED:
-        return None, "⚠️ Model is still loading or failed to load. Please refresh the page and try again."
-    
-    if image is None:
+    if input_image is None:
         return None, "Please upload an image first"
     
     try:
-        # Convert to PIL
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image).convert("RGB")
+        # Move model to GPU (ZeroGPU allocates it now)
+        pipe.to("cuda")
+        
+        # Convert to PIL if needed
+        if isinstance(input_image, np.ndarray):
+            image = Image.fromarray(input_image).convert("RGB")
         else:
-            image = image.convert("RGB")
+            image = input_image.convert("RGB")
         
         # Store original size
         original_size = image.size
         
-        # Resize for processing (max 512x512 for stability)
-        if max(image.size) > 512:
-            image.thumbnail((512, 512), Image.Resampling.LANCZOS)
+        # Resize for processing
+        max_size = 512
+        if max(image.size) > max_size:
+            ratio = max_size / max(image.size)
+            new_size = tuple(int(dim * ratio) for dim in image.size)
+            image = image.resize(new_size, Image.Resampling.LANCZOS)
         
         # Create mask
-        mask = create_mask(image)
+        mask = create_body_mask(image)
         
-        # Generate
-        prompt = CLOTHES[clothing_type]
-        negative_prompt = "nude, naked, nsfw, blurry, bad quality"
+        # Get prompt
+        prompt = CLOTHING_PROMPTS[clothing_type]
+        negative_prompt = "nude, naked, nsfw, bad quality, blurry, distorted"
         
-        with torch.no_grad():
-            if device == "cuda":
-                with torch.autocast("cuda"):
-                    result = pipe(
-                        prompt=prompt,
-                        negative_prompt=negative_prompt,
-                        image=image,
-                        mask_image=mask,
-                        num_inference_steps=steps,
-                        guidance_scale=7.5,
-                        strength=0.95
-                    ).images[0]
-            else:
-                result = pipe(
-                    prompt=prompt,
-                    negative_prompt=negative_prompt,
-                    image=image,
-                    mask_image=mask,
-                    num_inference_steps=steps,
-                    guidance_scale=7.5,
-                    strength=0.95
-                ).images[0]
+        # Generate with GPU
+        with torch.autocast("cuda"):
+            result = pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                image=image,
+                mask_image=mask,
+                num_inference_steps=num_steps,
+                guidance_scale=guidance_scale,
+                strength=0.95
+            ).images[0]
         
-        # Resize back to original
+        # Resize back
         if result.size != original_size:
             result = result.resize(original_size, Image.Resampling.LANCZOS)
         
+        # Move model back to CPU to free GPU
+        pipe.to("cpu")
+        torch.cuda.empty_cache()
+        
         return result, f"✅ Successfully added {clothing_type}!"
         
     except Exception as e:
         print(f"Generation error: {e}")
-        return None, f"❌ Error during generation: {str(e)}"
+        return None, f"Error: {str(e)}"
 
-# Create UI
-with gr.Blocks(title="Traditional Clothing AI") as app:
-    gr.Markdown(f"""
+# Create interface
+with gr.Blocks(title="Traditional Clothing AI - ZeroGPU", theme=gr.themes.Soft()) as app:
+    gr.Markdown("""
     # 👘 Traditional Clothing Addition Tool
     
-    **Device:** {device.upper()} {"🚀 Fast" if device == "cuda" else "🐌 Slow"}  
-    **Model:** {"✅ Ready" if MODEL_LOADED else "❌ Loading failed - please refresh"}
+    **Powered by ZeroGPU** 🚀 - Free GPU acceleration!
     
-    Transform your photos with traditional clothing from different cultures!
+    Add beautiful traditional clothing from various cultures to your photos.
+    Generation takes about 30-45 seconds per image.
     """)
     
     with gr.Row():
         with gr.Column():
-            input_img = gr.Image(
+            input_image = gr.Image(
                 label="Upload Your Photo",
                 type="pil"
             )
             
-            clothing_dropdown = gr.Dropdown(
-                choices=list(CLOTHES.keys()),
+            clothing_type = gr.Dropdown(
+                choices=list(CLOTHING_PROMPTS.keys()),
                 value="Indian Sari",
-                label="Choose Traditional Clothing"
+                label="Select Traditional Clothing"
             )
             
-            steps_slider = gr.Slider(
-                minimum=10,
-                maximum=50,
-                value=20,
-                step=5,
-                label="Quality (more steps = better but slower)"
-            )
+            with gr.Accordion("Advanced Settings", open=False):
+                num_steps = gr.Slider(
+                    minimum=15,
+                    maximum=50,
+                    value=25,
+                    step=5,
+                    label="Quality Steps",
+                    info="More steps = better quality but slower"
+                )
+                
+                guidance_scale = gr.Slider(
+                    minimum=5,
+                    maximum=15,
+                    value=7.5,
+                    step=0.5,
+                    label="Guidance Scale",
+                    info="Higher = more adherence to prompt"
+                )
             
             generate_btn = gr.Button(
                 "🎨 Add Traditional Clothing",
@@ -189,33 +157,38 @@ with gr.Blocks(title="Traditional Clothing AI") as app:
             )
         
         with gr.Column():
-            output_img = gr.Image(
+            output_image = gr.Image(
                 label="Result"
             )
             
             status_text = gr.Textbox(
                 label="Status",
-                interactive=False
+                placeholder="Upload an image and click generate..."
             )
     
-    # Examples (create dummy examples or remove if no images)
     gr.Markdown("""
-    ### 💡 Tips:
+    ---
+    ### 💡 Tips for Best Results:
     - Use clear, front-facing photos
     - Good lighting improves results
-    - Processing takes 15-30 seconds on GPU, 1-2 minutes on CPU
-    - Try different clothing types for variety
+    - The person should be fully visible
+    - Processing uses free GPU via ZeroGPU
     
     ### 🌍 Cultural Note:
-    This AI creates artistic interpretations of traditional clothing. 
-    Results may not be culturally accurate - please use respectfully.
+    This tool celebrates cultural diversity through traditional clothing.
+    AI-generated results are artistic interpretations.
+    Please use respectfully.
+    
+    ### ⚡ About ZeroGPU:
+    This Space uses Hugging Face's free ZeroGPU feature.
+    GPU is allocated only during generation, which saves resources!
     """)
     
     # Connect button
     generate_btn.click(
-        fn=process_image,
-        inputs=[input_img, clothing_dropdown, steps_slider],
-        outputs=[output_img, status_text]
+        fn=generate_clothing,
+        inputs=[input_image, clothing_type, num_steps, guidance_scale],
+        outputs=[output_image, status_text]
     )
 
 if __name__ == "__main__":