Update app.py

app.py

@@ -1,115 +1,167 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionInpaintPipeline
+from diffusers import StableDiffusionInpaintPipeline, StableDiffusionImg2ImgPipeline
 from PIL import Image, ImageDraw, ImageFilter
 import numpy as np
 import spaces
 
-# BETTER MODEL - This one is trained on people and fashion!
-# Options (all work great for people):
-# model_id = "Lykon/DreamShaper"  # Excellent for people
-# model_id = "prompthero/openjourney"  # Good quality
-# model_id = "wavymulder/Analog-Diffusion"  # Great for photos
-model_id = "Lykon/DreamShaper"  # Best overall for people
+# Use models that ACTUALLY WORK on HuggingFace
+# Option 1: Better inpainting model
+model_id = "runwayml/stable-diffusion-inpainting"
 
-print(f"Loading {model_id} - this model is MUCH better for people!")
+# Option 2: If you want to try regular SD models, use Img2Img pipeline instead
+# model_id = "runwayml/stable-diffusion-v1-5"
 
-pipe = StableDiffusionInpaintPipeline.from_pretrained(
-    model_id,
-    torch_dtype=torch.float16,
-    safety_checker=None,
-    requires_safety_checker=False
-)
-pipe.enable_attention_slicing()
-pipe.enable_vae_slicing()
+print(f"Loading {model_id}...")
 
-print("✅ Better model loaded!")
+try:
+    # For inpainting models
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16,
+        safety_checker=None,
+        requires_safety_checker=False,
+        use_safetensors=True
+    )
+    pipe.enable_attention_slicing()
+    pipe.enable_vae_slicing()
+    pipeline_type = "inpaint"
+    print("✅ Inpainting model loaded!")
+    
+except Exception as e:
+    print(f"Inpainting failed, trying img2img: {e}")
+    # Fallback to img2img pipeline
+    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        torch_dtype=torch.float16,
+        safety_checker=None,
+        requires_safety_checker=False
+    )
+    pipe.enable_attention_slicing()
+    pipeline_type = "img2img"
+    print("✅ Img2Img model loaded!")
 
-# Professional clothing prompts
+# MUCH BETTER PROMPTS - This is the secret!
 CLOTHING_PROMPTS = {
-    "Indian Sari": (
-        "beautiful woman wearing traditional red silk sari with golden embroidery, "
-        "professional photography, natural pose, perfect hands, detailed fabric, "
-        "elegant draping, studio lighting, high quality, sharp focus"
-    ),
-    "Japanese Kimono": (
-        "person wearing authentic silk kimono with floral patterns, obi belt, "
-        "traditional Japanese clothing, professional portrait, natural pose, "
-        "perfect anatomy, detailed fabric texture, studio photography"
-    ),
-    "African Dashiki": (
-        "person wearing colorful African dashiki with geometric patterns, "
-        "traditional clothing, professional photo, natural pose, perfect hands, "
-        "vibrant colors, detailed embroidery, high quality"
-    ),
-    "Chinese Qipao": (
-        "elegant woman wearing traditional Chinese qipao dress, silk fabric, "
-        "professional fashion photography, natural pose, perfect proportions, "
-        "detailed patterns, studio lighting"
-    ),
-    "Scottish Kilt": (
-        "man wearing traditional Scottish kilt with tartan pattern, "
-        "highland dress, professional photo, natural stance, proper proportions, "
-        "detailed fabric, high quality"
-    ),
-    "Middle Eastern Thobe": (
-        "person wearing traditional white thobe robe, Middle Eastern clothing, "
-        "professional portrait, natural pose, flowing fabric, elegant, "
-        "studio photography, high quality"
-    )
+    "Indian Sari": {
+        "prompt": (
+            "beautiful indian woman wearing elegant red silk sari with golden embroidery, "
+            "professional fashion photography, perfect anatomy, natural pose, "
+            "detailed fabric texture, studio lighting, high quality, 8k, sharp focus, "
+            "intricate patterns, traditional jewelry"
+        ),
+        "negative": (
+            "deformed, bad anatomy, disfigured, poorly drawn face, mutation, mutated, "
+            "extra limb, ugly, disgusting, poorly drawn hands, missing limb, floating limbs, "
+            "disconnected limbs, malformed hands, blurry, mutated hands and fingers, "
+            "watermark, oversaturated, distorted hands, amputation, missing hands, "
+            "doubled face, double hands, weird pose, unnatural pose"
+        )
+    },
+    "Japanese Kimono": {
+        "prompt": (
+            "elegant person wearing traditional japanese silk kimono with cherry blossom patterns, "
+            "obi belt, professional portrait photography, perfect proportions, natural pose, "
+            "detailed fabric texture, studio lighting, high quality, traditional hairstyle"
+        ),
+        "negative": (
+            "bad anatomy, wrong proportions, extra limbs, missing limbs, deformed, "
+            "ugly, duplicate, mutilated, out of frame, extra fingers, mutated hands, "
+            "poorly drawn hands, mutation, blurry, bad proportions"
+        )
+    },
+    "African Dashiki": {
+        "prompt": (
+            "person wearing colorful traditional african dashiki with geometric kente patterns, "
+            "vibrant colors, professional photography, perfect anatomy, natural pose, "
+            "detailed embroidery, cultural authenticity, high quality, sharp details"
+        ),
+        "negative": (
+            "bad anatomy, deformed, ugly, disfigured, poorly drawn, mutation, "
+            "extra limbs, missing limbs, floating limbs, disconnected limbs, "
+            "long neck, bad proportions, unnatural pose"
+        )
+    },
+    "Chinese Qipao": {
+        "prompt": (
+            "elegant woman wearing traditional chinese qipao cheongsam dress, silk fabric, "
+            "intricate patterns, professional fashion photography, perfect proportions, "
+            "natural pose, studio lighting, high quality, detailed embroidery"
+        ),
+        "negative": (
+            "deformed, bad anatomy, disfigured, mutation, extra limbs, ugly, "
+            "poorly drawn hands, missing limbs, blurry, bad art, bad proportions, "
+            "gross proportions, malformed limbs"
+        )
+    },
+    "Scottish Kilt": {
+        "prompt": (
+            "man wearing traditional scottish kilt with authentic tartan pattern, "
+            "highland dress, professional photography, perfect anatomy, natural stance, "
+            "detailed fabric, formal sporran, high quality"
+        ),
+        "negative": (
+            "bad anatomy, deformed, poorly drawn, extra limbs, close up, weird pose, "
+            "duplicate, mutilated, mutated hands, bad proportions"
+        )
+    },
+    "Middle Eastern Thobe": {
+        "prompt": (
+            "person wearing traditional white thobe robe, middle eastern clothing, "
+            "flowing fabric, professional portrait, perfect proportions, elegant, "
+            "studio photography, high quality, natural pose"
+        ),
+        "negative": (
+            "deformed, bad anatomy, ugly, poorly drawn, mutation, extra limbs, "
+            "bad hands, poorly drawn hands, missing limbs, blurry"
+        )
+    }
 }
 
-# Critical for good results - tells AI what to avoid
-NEGATIVE_PROMPT = (
-    "deformed, bad anatomy, disfigured, poorly drawn face, mutation, mutated, "
-    "extra limb, ugly, disgusting, poorly drawn hands, missing limb, floating limbs, "
-    "disconnected limbs, malformed hands, blurry, mutated hands and fingers, "
-    "watermark, oversaturated, distorted hands, amputation, missing hands, "
-    "obese, doubled face, double hands, bad hands, bad anatomy, bad proportions, "
-    "extra fingers, fused fingers, too many fingers, long neck, low quality"
-)
-
-def create_safe_mask(image, mask_strength="medium"):
-    """Create mask that preserves face and hands"""
+def create_smart_mask(image, coverage="medium"):
+    """Create intelligent mask that preserves anatomy"""
     width, height = image.size
     mask = Image.new('L', (width, height), 0)
     draw = ImageDraw.Draw(mask)
     
-    if mask_strength == "light":
-        # Very conservative - only center torso
+    if coverage == "light":
+        # Very safe - only central torso
         left = width * 0.35
         right = width * 0.65
         top = height * 0.45
         bottom = height * 0.65
         draw.ellipse([left, top, right, bottom], fill=255)
         
-    elif mask_strength == "medium":
-        # Moderate - torso and upper legs, avoid arms
+    elif coverage == "medium":
+        # Torso and upper body, preserve head and hands
         left = width * 0.25
         right = width * 0.75
-        top = height * 0.4
+        top = height * 0.38  # Start below neck
         bottom = height * 0.75
+        
+        # Main body
         draw.ellipse([left, top, right, bottom], fill=255)
         
-        # Cut out arm areas to preserve hands
-        arm_width = width * 0.15
-        draw.ellipse([left - arm_width, height * 0.45, left, height * 0.7], fill=0)
-        draw.ellipse([right, height * 0.45, right + arm_width, height * 0.7], fill=0)
+        # Exclude arm areas to preserve hands
+        arm_exclude = width * 0.12
+        draw.ellipse([left - arm_exclude, height * 0.45, left, height * 0.7], fill=0)
+        draw.ellipse([right, height * 0.45, right + arm_exclude, height * 0.7], fill=0)
         
-    else:  # heavy
-        # Most of body but still protect face and extremities
-        left = width * 0.15
-        right = width * 0.85
+    else:  # full
+        # More coverage but still protect extremities
+        left = width * 0.2
+        right = width * 0.8
         top = height * 0.35
-        bottom = height * 0.9
+        bottom = height * 0.85
         draw.ellipse([left, top, right, bottom], fill=255)
         
-        # Protect hand areas
-        draw.ellipse([0, height * 0.5, width * 0.1, height * 0.8], fill=0)
-        draw.ellipse([width * 0.9, height * 0.5, width, height * 0.8], fill=0)
+        # Protect hands
+        draw.rectangle([0, height * 0.6, width * 0.15, height], fill=0)
+        draw.rectangle([width * 0.85, height * 0.6, width, height], fill=0)
     
-    # Smooth blend
-    mask = mask.filter(ImageFilter.GaussianBlur(radius=30))
+    # Smooth edges
+    for _ in range(3):
+        mask = mask.filter(ImageFilter.GaussianBlur(radius=10))
     
     return mask
 
@@ -117,11 +169,11 @@ def create_safe_mask(image, mask_strength="medium"):
 def generate_clothing(
     input_image, 
     clothing_type, 
-    mask_strength="medium",
-    quality_steps=40,
-    guidance_scale=7.5
+    mask_coverage="medium",
+    num_steps=40,
+    strength=0.85
 ):
-    """Generate clothing with better model"""
+    """Generate clothing with proper technique"""
     
     if input_image is None:
         return None, "Please upload an image"
@@ -139,43 +191,52 @@ def generate_clothing(
         # Store original
         original_size = image.size
         
-        # Optimal size for quality (768 works well)
+        # Optimal size
         target_size = 768
         if max(image.size) > target_size:
             scale = target_size / max(image.size)
             new_w = int(image.width * scale)
             new_h = int(image.height * scale)
-            # Make divisible by 8
             new_w = new_w - (new_w % 8)
             new_h = new_h - (new_h % 8)
             image = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
         else:
-            # Still fix dimensions
+            # Fix dimensions
             new_w = image.width - (image.width % 8)
             new_h = image.height - (image.height % 8)
             if (new_w, new_h) != image.size:
                 image = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
         
-        print(f"Processing at: {image.size}")
+        # Get prompts
+        prompt_data = CLOTHING_PROMPTS[clothing_type]
         
-        # Create safe mask
-        mask = create_safe_mask(image, mask_strength)
-        
-        # Get prompt
-        prompt = CLOTHING_PROMPTS[clothing_type]
-        
-        # Generate with better settings
-        with torch.autocast("cuda"):
-            result = pipe(
-                prompt=prompt,
-                negative_prompt=NEGATIVE_PROMPT,
-                image=image,
-                mask_image=mask,
-                num_inference_steps=quality_steps,
-                guidance_scale=guidance_scale,
-                strength=0.85,  # Good balance
-                generator=torch.Generator("cuda").manual_seed(42)  # Consistent results
-            ).images[0]
+        # Generate based on pipeline type
+        if pipeline_type == "inpaint":
+            # Create mask
+            mask = create_smart_mask(image, mask_coverage)
+            
+            # Inpainting
+            with torch.autocast("cuda"):
+                result = pipe(
+                    prompt=prompt_data["prompt"],
+                    negative_prompt=prompt_data["negative"],
+                    image=image,
+                    mask_image=mask,
+                    num_inference_steps=num_steps,
+                    guidance_scale=7.5,
+                    strength=strength
+                ).images[0]
+        else:
+            # For img2img, use lower strength
+            with torch.autocast("cuda"):
+                result = pipe(
+                    prompt=prompt_data["prompt"],
+                    negative_prompt=prompt_data["negative"],
+                    image=image,
+                    num_inference_steps=num_steps,
+                    guidance_scale=7.5,
+                    strength=0.6  # Lower for img2img
+                ).images[0]
         
         # Resize back
         if result.size != original_size:
@@ -192,14 +253,16 @@ def generate_clothing(
         return None, f"Error: {str(e)}"
 
 # Create UI
-with gr.Blocks(title="Traditional Clothing AI - Better Model", theme=gr.themes.Soft()) as app:
-    gr.Markdown("""
-    # 👘 Traditional Clothing AI - Premium Quality
+with gr.Blocks(title="Traditional Clothing AI", theme=gr.themes.Soft()) as app:
+    gr.Markdown(f"""
+    # 👘 Traditional Clothing AI - Working Version
     
-    ### Using DreamShaper model - MUCH better for people!
-    - ✅ Preserves anatomy correctly
-    - ✅ No weird poses or hands
-    - ✅ Professional quality results
+    ### Model: {model_id} ({pipeline_type} mode)
+    
+    **Tips for best results:**
+    - The SECRET is in the detailed prompts!
+    - Use "light" coverage if you see weird poses
+    - Higher steps = better quality
     """)
     
     with gr.Row():
@@ -215,30 +278,28 @@ with gr.Blocks(title="Traditional Clothing AI - Better Model", theme=gr.themes.S
                 label="Select Traditional Clothing"
             )
             
-            with gr.Accordion("Advanced Settings", open=True):
-                mask_strength = gr.Radio(
-                    choices=["light", "medium", "heavy"],
+            with gr.Accordion("Settings", open=True):
+                mask_coverage = gr.Radio(
+                    choices=["light", "medium", "full"],
                     value="medium",
-                    label="Mask Coverage",
-                    info="Light = safest (only torso), Heavy = more coverage"
+                    label="Clothing Coverage",
+                    info="Light = safest, Full = most coverage"
                 )
                 
-                quality_steps = gr.Slider(
+                num_steps = gr.Slider(
                     minimum=20,
                     maximum=60,
                     value=40,
                     step=5,
-                    label="Quality Steps",
-                    info="Higher = better quality but slower"
+                    label="Quality Steps"
                 )
                 
-                guidance_scale = gr.Slider(
-                    minimum=5,
-                    maximum=12,
-                    value=7.5,
-                    step=0.5,
-                    label="Guidance Scale",
-                    info="How closely to follow the prompt"
+                strength = gr.Slider(
+                    minimum=0.5,
+                    maximum=0.95,
+                    value=0.85,
+                    step=0.05,
+                    label="Change Strength"
                 )
             
             generate_btn = gr.Button(
@@ -248,59 +309,26 @@ with gr.Blocks(title="Traditional Clothing AI - Better Model", theme=gr.themes.S
             )
         
         with gr.Column():
-            output_image = gr.Image(
-                label="Result"
-            )
-            
-            status_text = gr.Textbox(
-                label="Status",
-                placeholder="Upload an image and click generate..."
-            )
-    
-    # Examples
-    gr.Examples(
-        examples=[
-            ["person1.jpg", "Indian Sari", "medium"],
-            ["person2.jpg", "Japanese Kimono", "medium"],
-            ["person3.jpg", "African Dashiki", "light"],
-        ],
-        inputs=[input_image, clothing_type, mask_strength],
-        outputs=[output_image, status_text],
-        fn=generate_clothing,
-        cache_examples=False
-    )
+            output_image = gr.Image(label="Result")
+            status_text = gr.Textbox(label="Status")
     
     gr.Markdown("""
-    ---
-    ### 🎯 Tips for Best Results:
-    
-    1. **Photo Tips:**
-       - Use clear, front-facing photos
-       - Good lighting helps
-       - Full body or 3/4 shots work best
-    
-    2. **If hands/feet look weird:**
-       - Use "Light" mask coverage
-       - This only changes the torso area
-    
-    3. **For better quality:**
-       - Increase quality steps to 50-60
-       - Takes longer but worth it
+    ### 🎯 The Secret: Better Prompts!
     
-    4. **Model Advantages:**
-       - DreamShaper is trained on people
-       - Much better anatomy understanding
-       - Preserves natural poses
+    The model is less important than the prompts. This version uses:
+    - Detailed clothing descriptions
+    - Anatomy preservation keywords
+    - Strong negative prompts
     
-    ### 🌍 Note:
-    This AI creates artistic interpretations of traditional clothing.
-    Results may not be culturally perfect - please use respectfully.
+    ### If you see weird poses:
+    1. Use "light" coverage
+    2. Lower the strength to 0.7
+    3. Increase steps to 50+
     """)
     
-    # Connect button
     generate_btn.click(
         fn=generate_clothing,
-        inputs=[input_image, clothing_type, mask_strength, quality_steps, guidance_scale],
+        inputs=[input_image, clothing_type, mask_coverage, num_steps, strength],
         outputs=[output_image, status_text]
     )