Update app.py

app.py

@@ -5,190 +5,217 @@ from PIL import Image, ImageDraw, ImageFilter
 import numpy as np
 import os
 
-# Set memory allocation
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
+# Clear cache if needed
+os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = '1'
 
-# Check device - wait for GPU
+# Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Device detected: {device}")
+dtype = torch.float16 if device == "cuda" else torch.float32
+print(f"Using device: {device}")
 
-if device == "cuda":
-    print(f"GPU: {torch.cuda.get_device_name(0)}")
-    print(f"Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
+# Try loading model with better error handling
+MODEL_LOADED = False
+pipe = None
 
-# Initialize model
 try:
-    print("Loading Stable Diffusion Inpainting model...")
-    
-    model_id = "runwayml/stable-diffusion-inpainting"
+    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(
-        model_id,
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        "runwayml/stable-diffusion-inpainting",
         revision="fp16" if device == "cuda" else "main",
+        torch_dtype=dtype,
         safety_checker=None,
-        requires_safety_checker=False
+        requires_safety_checker=False,
+        use_safetensors=True,  # Force safetensors format
+        local_files_only=False  # Allow downloading
     )
     
-    # Move to GPU if available, otherwise keep on CPU
-    if device == "cuda":
-        pipe = pipe.to(device)
-        pipe.enable_attention_slicing()
-        print("Model loaded on GPU!")
-    else:
-        # Don't use CPU offload on HF Spaces - it causes issues
-        print("Warning: Running on CPU - will be slow!")
+    pipe = pipe.to(device)
+    pipe.enable_attention_slicing()
     
     MODEL_LOADED = True
+    print("✅ Model loaded successfully!")
     
 except Exception as e:
-    print(f"Error loading model: {e}")
-    MODEL_LOADED = False
-    pipe = None
+    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
 
 # Clothing prompts
-CLOTHING_PROMPTS = {
-    "Indian Sari": "woman wearing elegant red and gold sari, traditional Indian saree dress",
-    "Japanese Kimono": "person wearing beautiful kimono with floral patterns, traditional Japanese clothing",
-    "African Dashiki": "person wearing colorful dashiki with patterns, traditional African clothing",
+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 Scottish kilt with tartan pattern, traditional highland dress",
+    "Scottish Kilt": "man wearing traditional Scottish kilt with tartan pattern",
     "Middle Eastern Thobe": "person wearing white thobe robe, traditional Middle Eastern clothing"
 }
 
-def create_body_mask(image):
-    """Create mask for clothing area"""
-    width, height = image.size
-    mask = Image.new('L', (width, height), 0)
+def create_mask(image):
+    """Create body mask"""
+    w, h = image.size
+    mask = Image.new('L', (w, h), 0)
     draw = ImageDraw.Draw(mask)
     
-    # Body area
-    left = width * 0.2
-    top = height * 0.25
-    right = width * 0.8
-    bottom = height * 0.95
+    # Body area ellipse
+    left = w * 0.2
+    top = h * 0.25
+    right = w * 0.8
+    bottom = h * 0.9
     
     draw.ellipse([left, top, right, bottom], fill=255)
-    mask = mask.filter(ImageFilter.GaussianBlur(radius=15))
+    mask = mask.filter(ImageFilter.GaussianBlur(radius=20))
     
     return mask
 
-def add_traditional_clothing(input_image, clothing_type, num_steps=20, guidance_scale=7.5):
-    """Add traditional clothing to image"""
-    
-    if input_image is None:
-        return None, "Please upload an image"
+def process_image(image, clothing_type, steps=20):
+    """Generate traditional clothing"""
     
     if not MODEL_LOADED:
-        return None, "Model failed to load. Please refresh the page."
+        return None, "⚠️ Model is still loading or failed to load. Please refresh the page and try again."
+    
+    if image is None:
+        return None, "Please upload an image first"
     
     try:
         # Convert to PIL
-        if isinstance(input_image, np.ndarray):
-            image = Image.fromarray(input_image).convert("RGB")
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image).convert("RGB")
         else:
-            image = input_image.convert("RGB")
+            image = image.convert("RGB")
         
         # Store original size
         original_size = image.size
         
-        # 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)
+        # Resize for processing (max 512x512 for stability)
+        if max(image.size) > 512:
+            image.thumbnail((512, 512), Image.Resampling.LANCZOS)
         
         # Create mask
-        mask = create_body_mask(image)
-        
-        # Get prompt
-        prompt = CLOTHING_PROMPTS[clothing_type]
-        negative_prompt = "nude, naked, nsfw, bad quality, blurry"
+        mask = create_mask(image)
         
         # Generate
-        if device == "cuda":
-            with torch.autocast("cuda"):
+        prompt = CLOTHES[clothing_type]
+        negative_prompt = "nude, naked, nsfw, blurry, bad quality"
+        
+        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=num_steps,
-                    guidance_scale=guidance_scale,
+                    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=num_steps,
-                guidance_scale=guidance_scale,
-                strength=0.95
-            ).images[0]
         
-        # Resize back
+        # Resize back to original
         if result.size != original_size:
             result = result.resize(original_size, Image.Resampling.LANCZOS)
         
-        return result, f"Successfully added {clothing_type}!"
+        return result, f"✅ Successfully added {clothing_type}!"
         
     except Exception as e:
         print(f"Generation error: {e}")
-        return None, f"Error: {str(e)}"
+        return None, f"❌ Error during generation: {str(e)}"
 
-# Create interface
-with gr.Blocks(title="Traditional Clothing Addition") as app:
+# Create UI
+with gr.Blocks(title="Traditional Clothing AI") as app:
     gr.Markdown(f"""
     # 👘 Traditional Clothing Addition Tool
     
-    **Status:** Running on {device.upper()} {"✅" if device == "cuda" else "⚠️ (Slow)"}
+    **Device:** {device.upper()} {"🚀 Fast" if device == "cuda" else "🐌 Slow"}  
+    **Model:** {"✅ Ready" if MODEL_LOADED else "❌ Loading failed - please refresh"}
     
-    Add beautiful traditional clothing to your photos using AI.
+    Transform your photos with traditional clothing from different cultures!
     """)
     
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(label="Upload Photo", type="pil")
+            input_img = gr.Image(
+                label="Upload Your Photo",
+                type="pil"
+            )
             
-            clothing_type = gr.Dropdown(
-                choices=list(CLOTHING_PROMPTS.keys()),
+            clothing_dropdown = gr.Dropdown(
+                choices=list(CLOTHES.keys()),
                 value="Indian Sari",
-                label="Select Traditional Clothing"
+                label="Choose Traditional Clothing"
             )
             
-            with gr.Accordion("Settings", open=False):
-                num_steps = gr.Slider(
-                    minimum=10, maximum=50, value=20, step=5,
-                    label="Steps (more = better quality)"
-                )
-                guidance_scale = gr.Slider(
-                    minimum=5, maximum=15, value=7.5, step=0.5,
-                    label="Guidance Scale"
-                )
+            steps_slider = gr.Slider(
+                minimum=10,
+                maximum=50,
+                value=20,
+                step=5,
+                label="Quality (more steps = better but slower)"
+            )
             
-            generate_btn = gr.Button("🎨 Add Traditional Clothing", variant="primary")
+            generate_btn = gr.Button(
+                "🎨 Add Traditional Clothing",
+                variant="primary",
+                size="lg"
+            )
         
         with gr.Column():
-            output_image = gr.Image(label="Result")
-            status = gr.Textbox(label="Status")
+            output_img = gr.Image(
+                label="Result"
+            )
+            
+            status_text = gr.Textbox(
+                label="Status",
+                interactive=False
+            )
     
+    # Examples (create dummy examples or remove if no images)
     gr.Markdown("""
-    ### Tips:
-    - Processing takes 15-30 seconds on GPU
+    ### 💡 Tips:
     - Use clear, front-facing photos
-    - Try different settings for variety
+    - Good lighting improves results
+    - Processing takes 15-30 seconds on GPU, 1-2 minutes on CPU
+    - Try different clothing types for variety
     
-    ### Note:
-    AI-generated traditional clothing may not be culturally accurate. Use respectfully.
+    ### 🌍 Cultural Note:
+    This AI creates artistic interpretations of traditional clothing. 
+    Results may not be culturally accurate - please use respectfully.
     """)
     
+    # Connect button
     generate_btn.click(
-        add_traditional_clothing,
-        inputs=[input_image, clothing_type, num_steps, guidance_scale],
-        outputs=[output_image, status]
+        fn=process_image,
+        inputs=[input_img, clothing_dropdown, steps_slider],
+        outputs=[output_img, status_text]
     )
 
 if __name__ == "__main__":