import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image, ImageDraw, ImageFilter, ImageOps
import numpy as np
import spaces

# Use regular SD pipeline - more reliable than inpainting
model_id = "runwayml/stable-diffusion-v1-5"

pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    safety_checker=None,
    requires_safety_checker=False
)
pipe.enable_attention_slicing()

# We'll simulate inpainting using img2img + compositing
PROMPTS = {
    "Sari": "woman wearing beautiful red and gold traditional indian sari, professional fashion photo",
    "Kimono": "person wearing elegant traditional japanese kimono, professional portrait",
    "Dashiki": "person wearing vibrant african dashiki with patterns, professional photo",
    "Qipao": "woman wearing traditional chinese qipao dress, elegant, professional"
}

def create_clothing_area_mask(image):
    """Create mask for clothing area only"""
    w, h = image.size
    
    # Create gradient mask - stronger in center/torso
    mask = Image.new('L', (w, h), 0)
    draw = ImageDraw.Draw(mask)
    
    # Torso area gets full opacity
    torso_area = [w*0.25, h*0.35, w*0.75, h*0.75]
    draw.ellipse(torso_area, fill=255)
    
    # Fade out at edges
    mask = mask.filter(ImageFilter.GaussianBlur(radius=40))
    
    return mask

@spaces.GPU(duration=60)
def generate_clothing(image, clothing_type):
    if image is None:
        return None, "Please upload an image"
    
    try:
        pipe.to("cuda")
        
        # Convert image
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image).convert("RGB")
        
        original = image.copy()
        
        # Resize for processing
        if max(image.size) > 512:
            image.thumbnail((512, 512), Image.Resampling.LANCZOS)
            original.thumbnail((512, 512), Image.Resampling.LANCZOS)
        
        # Fix dimensions
        w, h = image.size
        w = w - (w % 8)
        h = h - (h % 8)
        image = image.resize((w, h), Image.Resampling.LANCZOS)
        original = original.resize((w, h), Image.Resampling.LANCZOS)
        
        # Generate new image with clothing
        prompt = PROMPTS[clothing_type] + ", high quality, professional photography"
        negative = "ugly, deformed, bad anatomy, bad hands"
        
        # Use image-to-image generation
        with torch.autocast("cuda"):
            generated = pipe(
                prompt=prompt,
                negative_prompt=negative,
                image=image,
                strength=0.7,  # Moderate transformation
                num_inference_steps=30,
                guidance_scale=7.5
            ).images[0]
        
        # Create smooth blend using mask
        mask = create_clothing_area_mask(original)
        
        # Composite: use generated for clothing area, original for face/hands
        final = Image.composite(generated, original, mask)
        
        pipe.to("cpu")
        torch.cuda.empty_cache()
        
        return final, f"✅ {clothing_type} added successfully!"
        
    except Exception as e:
        return None, f"Error: {str(e)}"

# UI
with gr.Blocks(title="Traditional Clothing AI") as app:
    gr.Markdown("""
    # 👘 Traditional Clothing AI - Alternative Method
    
    This version uses regular SD model + smart blending (avoids inpainting issues).
    """)
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Upload Photo")
            clothing = gr.Dropdown(list(PROMPTS.keys()), value="Sari", label="Clothing")
            btn = gr.Button("Generate", variant="primary")
        
        with gr.Column():
            output = gr.Image(label="Result")
            status = gr.Textbox(label="Status")
    
    gr.Markdown("""
    ### Why this works better:
    - Uses standard SD model (always downloads correctly)
    - Smart blending preserves face/hands
    - No special inpainting model needed
    """)
    
    btn.click(generate_clothing, [input_img, clothing], [output, status])

app.launch()