Update app.py

app.py

@@ -1,140 +1,57 @@
 import gradio as gr
 import torch
 from diffusers import StableDiffusionInpaintPipeline
-from PIL import Image, ImageDraw, ImageFilter, ImageEnhance
+from PIL import Image, ImageDraw, ImageFilter
 import numpy as np
 import spaces
 
 # Load model
-inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
     "stabilityai/stable-diffusion-2-inpainting",
     torch_dtype=torch.float16,
     safety_checker=None,
     requires_safety_checker=False
 )
-inpaint_pipe.enable_attention_slicing()
-inpaint_pipe.enable_vae_slicing()
-inpaint_pipe.enable_vae_tiling()
+pipe.enable_attention_slicing()
 
-print("✅ Model loaded!")
-
-CLOTHING_PROMPTS = {
-    "Indian Sari": "woman wearing luxurious red and gold silk sari with intricate embroidery, traditional Indian saree, professional fashion photography, studio lighting, ultra detailed fabric texture, 8k quality",
-    "Japanese Kimono": "person wearing exquisite silk kimono with cherry blossom patterns, traditional Japanese formal wear, professional portrait, studio lighting, highly detailed fabric, photorealistic",
-    "African Dashiki": "person wearing vibrant African dashiki with authentic kente patterns, traditional clothing, professional photography, rich colors, detailed textile work, high resolution",
-    "Chinese Qipao": "elegant woman in traditional Chinese qipao cheongsam, silk dress with intricate patterns, professional fashion shoot, studio lighting, ultra high quality",
-    "Scottish Kilt": "man wearing traditional Scottish highland kilt with tartan pattern, formal Scottish attire, professional photography, detailed fabric texture",
-    "Middle Eastern Thobe": "person wearing flowing white thobe robe, traditional Middle Eastern clothing, elegant fabric, studio portrait, high resolution"
+CLOTHES = {
+    "Indian Sari": "woman wearing beautiful red and gold sari, traditional Indian dress, high quality photo",
+    "Japanese Kimono": "person wearing elegant kimono, traditional Japanese clothing, professional photo",
+    "African Dashiki": "person wearing colorful dashiki, traditional African clothing, detailed",
+    "Chinese Qipao": "woman wearing elegant qipao dress, traditional Chinese clothing",
+    "Scottish Kilt": "man wearing Scottish kilt, traditional highland dress",
+    "Middle Eastern Thobe": "person wearing white thobe, traditional Middle Eastern clothing"
 }
 
-def make_divisible_by_8(image):
-    """Ensure image dimensions are divisible by 8"""
-    width, height = image.size
-    
-    # Calculate new dimensions divisible by 8
-    new_width = width - (width % 8)
-    new_height = height - (height % 8)
-    
-    # Only resize if needed
-    if new_width != width or new_height != height:
-        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
-    
-    return image
+def make_divisible_by_8(width, height):
+    """Ensure dimensions are divisible by 8"""
+    return width - (width % 8), height - (height % 8)
 
-def resize_with_aspect_ratio(image, target_size):
-    """Resize keeping aspect ratio and ensuring divisible by 8"""
-    width, height = image.size
-    
-    # Calculate scaling factor
-    scale = target_size / max(width, height)
-    
-    # Calculate new dimensions
-    new_width = int(width * scale)
-    new_height = int(height * scale)
-    
-    # Make divisible by 8
-    new_width = new_width - (new_width % 8)
-    new_height = new_height - (new_height % 8)
-    
-    # Ensure minimum size
-    new_width = max(new_width, 64)
-    new_height = max(new_height, 64)
-    
-    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
-
-def create_professional_mask(image, face_margin=0.35):
-    """Create mask avoiding face area"""
-    width, height = image.size
+def create_body_mask(image_size):
+    """Create mask for body area only"""
+    width, height = image_size
     mask = Image.new('L', (width, height), 0)
     draw = ImageDraw.Draw(mask)
     
-    # Face-safe area
-    face_bottom = height * face_margin
-    
-    # Body area
-    body_coords = [
-        width * 0.1,
-        face_bottom,
-        width * 0.9,
-        height * 0.98
-    ]
-    
-    # Draw body
-    draw.ellipse(body_coords, fill=255)
-    
-    # Gradient for smooth transition
-    for i in range(30):
-        opacity = int(255 * (i / 30))
-        y = face_bottom - (30 - i)
-        if y >= 0:
-            draw.rectangle([body_coords[0], y, body_coords[2], y + 1], fill=opacity)
+    # Body area (avoiding face)
+    top = height * 0.35  # Start below face
+    left = width * 0.1
+    right = width * 0.9
+    bottom = height * 0.98
     
-    # Smooth blur
+    draw.ellipse([left, top, right, bottom], fill=255)
     mask = mask.filter(ImageFilter.GaussianBlur(radius=25))
     
     return mask
 
-def enhance_image(image):
-    """Enhance image quality"""
-    # Sharpness
-    enhancer = ImageEnhance.Sharpness(image)
-    image = enhancer.enhance(1.2)
-    
-    # Color
-    enhancer = ImageEnhance.Color(image)
-    image = enhancer.enhance(1.1)
-    
-    return image
-
-def blend_images(original, generated, mask):
-    """Blend images smoothly"""
-    # Extra smooth blending
-    blend_mask = mask.filter(ImageFilter.GaussianBlur(radius=40))
-    
-    # Convert to RGBA
-    original_rgba = original.convert("RGBA")
-    generated_rgba = generated.convert("RGBA")
-    
-    # Composite
-    result = Image.composite(generated_rgba, original_rgba, blend_mask)
-    
-    return result.convert("RGB")
-
-@spaces.GPU(duration=120)
-def generate_professional(
-    input_image, 
-    clothing_type, 
-    face_margin=0.35,
-    quality_preset="ultra"
-):
-    """Generate with proper dimension handling"""
-    
+@spaces.GPU(duration=90)
+def generate_clothing(input_image, clothing_type, quality_mode="balanced"):
     if input_image is None:
         return None, "Please upload an image"
     
     try:
         # Move to GPU
-        inpaint_pipe.to("cuda")
+        pipe.to("cuda")
         
         # Convert to PIL
         if isinstance(input_image, np.ndarray):
@@ -142,118 +59,145 @@ def generate_professional(
         else:
             image = input_image.convert("RGB")
         
-        # Store original
-        original_image = image.copy()
+        # Store original size
         original_size = image.size
         
         # Quality settings
         quality_settings = {
-            "fast": {"size": 512, "steps": 30},
-            "balanced": {"size": 768, "steps": 50},
-            "ultra": {"size": 1024, "steps": 70}
+            "fast": {"size": 512, "steps": 25},
+            "balanced": {"size": 768, "steps": 40},
+            "ultra": {"size": 1024, "steps": 60}
         }
         
-        settings = quality_settings[quality_preset]
+        settings = quality_settings[quality_mode]
+        target_size = settings["size"]
         
-        # Resize with proper dimensions
-        if max(image.size) > settings["size"]:
-            image = resize_with_aspect_ratio(image, settings["size"])
-            original_resized = original_image.resize(image.size, Image.Resampling.LANCZOS)
+        # Calculate new size maintaining aspect ratio
+        if max(image.size) > target_size:
+            scale = target_size / max(image.size)
+            new_width = int(image.width * scale)
+            new_height = int(image.height * scale)
         else:
-            # Still ensure divisible by 8
-            image = make_divisible_by_8(image)
-            original_resized = original_image.resize(image.size, Image.Resampling.LANCZOS)
+            new_width = image.width
+            new_height = image.height
+        
+        # Make divisible by 8
+        new_width, new_height = make_divisible_by_8(new_width, new_height)
         
-        print(f"Processing at: {image.size} (divisible by 8)")
+        # Ensure minimum size
+        new_width = max(new_width, 64)
+        new_height = max(new_height, 64)
         
-        # Enhance
-        image = enhance_image(image)
+        # Resize all images to the same size
+        working_size = (new_width, new_height)
+        image_resized = image.resize(working_size, Image.Resampling.LANCZOS)
         
-        # Create mask
-        mask = create_professional_mask(image, face_margin)
+        # Create mask at the same size
+        mask = create_body_mask(working_size)
+        
+        print(f"Processing at size: {working_size}")
         
         # Generate
-        prompt = CLOTHING_PROMPTS[clothing_type]
-        negative_prompt = "blurry, low quality, distorted face, bad anatomy, ugly"
+        prompt = CLOTHES[clothing_type] + ", professional photography, preserve facial features"
+        negative_prompt = "blurry, low quality, distorted face, bad anatomy"
         
         with torch.autocast("cuda"):
-            result = inpaint_pipe(
+            result = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
-                image=image,
+                image=image_resized,
                 mask_image=mask,
                 num_inference_steps=settings["steps"],
-                guidance_scale=8.0,
-                strength=0.88
+                guidance_scale=7.5,
+                strength=0.85
             ).images[0]
         
-        # Blend
-        final = blend_images(original_resized, result, mask)
+        # Ensure result is the same size (it should be, but just in case)
+        if result.size != working_size:
+            result = result.resize(working_size, Image.Resampling.LANCZOS)
         
-        # Final enhancement
-        final = enhance_image(final)
+        # Blend with original to preserve face
+        # Create smooth blend mask
+        blend_mask = mask.filter(ImageFilter.GaussianBlur(radius=40))
         
-        # Resize back to original
+        # All images must be the same size for composite
+        assert image_resized.size == result.size == blend_mask.size, f"Size mismatch: {image_resized.size}, {result.size}, {blend_mask.size}"
+        
+        # Blend
+        final = Image.composite(result, image_resized, blend_mask)
+        
+        # Resize back to original size
         if final.size != original_size:
             final = final.resize(original_size, Image.Resampling.LANCZOS)
         
         # Cleanup
-        inpaint_pipe.to("cpu")
+        pipe.to("cpu")
         torch.cuda.empty_cache()
         
-        return final, f"✅ {clothing_type} applied successfully!"
+        return final, f"✅ Successfully added {clothing_type}!"
         
     except Exception as e:
+        print(f"Error details: {str(e)}")
         return None, f"Error: {str(e)}"
 
 # UI
-with gr.Blocks(title="Professional Clothing AI") as app:
+with gr.Blocks(title="Traditional Clothing AI", theme=gr.themes.Soft()) as app:
     gr.Markdown("""
-    # 👘 Professional Traditional Clothing AI
-    ### Perfect Face Preservation • Studio Quality
+    # 👘 Traditional Clothing AI - Face Preserved
+    
+    Add traditional clothing while keeping your face perfectly intact.
     """)
     
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(type="pil", label="Upload Photo")
+            input_image = gr.Image(
+                type="pil",
+                label="Upload Your Photo"
+            )
             
             clothing_type = gr.Dropdown(
-                choices=list(CLOTHING_PROMPTS.keys()),
+                choices=list(CLOTHES.keys()),
                 value="Indian Sari",
-                label="Traditional Clothing"
+                label="Select Traditional Clothing"
             )
             
-            with gr.Accordion("Settings", open=True):
-                face_margin = gr.Slider(
-                    0.25, 0.45, 0.35, 0.05,
-                    label="Face Protection Zone",
-                    info="Higher = more face area protected"
-                )
-                
-                quality_preset = gr.Radio(
-                    ["fast", "balanced", "ultra"],
-                    value="balanced",
-                    label="Quality",
-                    info="Ultra = best (2-3 min)"
-                )
+            quality_mode = gr.Radio(
+                choices=["fast", "balanced", "ultra"],
+                value="balanced",
+                label="Quality Mode",
+                info="Higher quality = longer processing time"
+            )
             
-            generate_btn = gr.Button("🎨 Generate", variant="primary", size="lg")
+            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_image = gr.Image(
+                label="Result"
+            )
+            
+            status_text = gr.Textbox(
+                label="Status",
+                placeholder="Upload an image and click generate..."
+            )
     
     gr.Markdown("""
-    ### Tips:
-    - Face margin 0.35 = perfect face preservation
-    - Balanced mode = good quality in ~1 minute
-    - Works with any image size
+    ### How it works:
+    - 🎯 Only modifies clothing area (below face)
+    - 😊 Your face remains untouched
+    - 🎨 Smooth blending for natural results
+    - ⚡ Fast mode: ~30 seconds
+    - 🔬 Ultra mode: ~2 minutes (best quality)
     """)
     
     generate_btn.click(
-        generate_professional,
-        inputs=[input_image, clothing_type, face_margin, quality_preset],
-        outputs=[output_image, status]
+        fn=generate_clothing,
+        inputs=[input_image, clothing_type, quality_mode],
+        outputs=[output_image, status_text]
     )
 
-app.launch()
+if __name__ == "__main__":
+    app.launch()