app.py

import gradio as gr
import numpy as np
import cv2
from PIL import Image, ImageDraw, ImageFont
import requests
import json
import random
import base64
import io
import os
from diffusers import StableDiffusionInpaintPipeline, StableDiffusionPipeline
import torch

# Initialize AI models
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe_inpaint = None
pipe_txt2img = None

# Base model photos (plain people to put outfits on)
MODEL_AVATARS = {
    "Fashion Model A": "https://i.postimg.cc/mD60gXJ6/male-model.png",
    "Fashion Model B": "https://i.postimg.cc/ZYWqF1bn/female-model.png",
    "Fashion Model C": "https://i.postimg.cc/Gh0bPtgg/female-model2.png",
    "Casual Model A": "https://i.postimg.cc/T2b6PG2V/male-model2.png"
}

def add_watermark(img):
    """Add watermark to the image"""
    if isinstance(img, np.ndarray):
        h, w, _ = img.shape
        img = cv2.putText(img, 'AI Outfit Designer', (int(0.02*w), h-10), 
                         cv2.FONT_HERSHEY_SIMPLEX, 0.6, (128, 128, 128), 1, cv2.LINE_AA)
    return img

def initialize_models():
    global pipe_inpaint, pipe_txt2img
    try:
        # Load inpainting model for outfit editing
        pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
            "runwayml/stable-diffusion-inpainting",
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            safety_checker=None,
            requires_safety_checker=False
        ).to(device)
        
        # Load text-to-image model for outfit generation
        pipe_txt2img = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            safety_checker=None,
            requires_safety_checker=False
        ).to(device)
        
        if device == "cuda":
            pipe_inpaint.enable_attention_slicing()
            pipe_txt2img.enable_attention_slicing()
            
        return True
    except Exception as e:
        print(f"Error loading models: {e}")
        return False

def create_placeholder_image(text="Model Image", size=(400, 600)):
    """Create placeholder image for models"""
    img = Image.new('RGB', size, color='lightgray')
    draw = ImageDraw.Draw(img)
    
    try:
        font = ImageFont.truetype("arial.ttf", 20)
    except:
        font = ImageFont.load_default()
    
    bbox = draw.textbbox((0, 0), text, font=font)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]
    
    x = (size[0] - text_width) // 2
    y = (size[1] - text_height) // 2
    
    draw.text((x, y), text, fill='black', font=font)
    return np.array(img)

def load_model_image_from_url(url):
    """Load model image from URL"""
    try:
        response = requests.get(url, timeout=10)
        if response.status_code == 200:
            image = Image.open(io.BytesIO(response.content))
            # Convert to RGB if needed
            if image.mode != 'RGB':
                image = image.convert('RGB')
            # Resize to standard size
            image = image.resize((400, 600), Image.Resampling.LANCZOS)
            return np.array(image)
        else:
            return create_placeholder_image("Failed to load model")
    except Exception as e:
        print(f"Error loading model image: {e}")
        return create_placeholder_image("Model loading error")

def get_outfit_tryon_result(model_image, garment_top, garment_bottom, style_prompt="", selected_model=None):
    """Generate outfit try-on result"""
    try:
        # Determine which model to use
        if model_image is not None:
            # Use uploaded image
            base_model = model_image
            model_source = "uploaded photo"
        elif selected_model and selected_model in MODEL_AVATARS:
            # Load selected model from URL
            base_model = load_model_image_from_url(MODEL_AVATARS[selected_model])
            model_source = selected_model
        else:
            return create_placeholder_image("Please upload a photo or select a model")
        
        # If no garments provided, generate based on style
        if garment_top is None and garment_bottom is None and not style_prompt:
            return create_placeholder_image("Please upload garments or describe a style")
        
        if garment_top is None and garment_bottom is None and style_prompt:
            # Generate outfit based on text description
            if pipe_txt2img:
                prompt = f"professional fashion photo of a person wearing {style_prompt}, high quality, detailed clothing, same person as reference"
                negative_prompt = "blurry, low quality, distorted, nude, different person"
                
                generated_image = pipe_txt2img(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    num_inference_steps=25,
                    guidance_scale=7.5,
                    width=400,
                    height=600
                ).images[0]
                
                result = np.array(generated_image)
            else:
                # Create styled placeholder showing the base model with text overlay
                result = create_styled_placeholder(base_model, style_prompt)
        else:
            # Apply uploaded garments to the base model
            result = simulate_tryon(base_model, garment_top, garment_bottom, model_source)
        
        # Add watermark
        final_result = add_watermark(result)
        return final_result
        
    except Exception as e:
        print(f"Error in outfit generation: {e}")
        return create_placeholder_image("Generation Error")

def create_styled_placeholder(base_model, style_text):
    """Create a styled placeholder showing base model with style text"""
    if isinstance(base_model, np.ndarray):
        result = base_model.copy()
    else:
        result = create_placeholder_image("Base Model")
    
    # Add style overlay
    img_pil = Image.fromarray(result)
    draw = ImageDraw.Draw(img_pil)
    
    try:
        font = ImageFont.truetype("arial.ttf", 18)
        small_font = ImageFont.truetype("arial.ttf", 14)
    except:
        font = ImageFont.load_default()
        small_font = ImageFont.load_default()
    
    # Add semi-transparent background for text
    overlay = Image.new('RGBA', img_pil.size, (0,0,0,0))
    overlay_draw = ImageDraw.Draw(overlay)
    overlay_draw.rectangle([10, 10, 390, 80], fill=(0,0,0,128))
    
    img_pil = Image.alpha_composite(img_pil.convert('RGBA'), overlay).convert('RGB')
    draw = ImageDraw.Draw(img_pil)
    
    # Add text
    draw.text((20, 20), "Generated Style:", fill='white', font=small_font)
    draw.text((20, 45), style_text[:40] + "..." if len(style_text) > 40 else style_text, 
              fill='yellow', font=font)
    
    return np.array(img_pil)

def simulate_tryon(model_img, garment_top, garment_bottom, model_source=""):
    """Simulate outfit try-on with base model and garments"""
    if isinstance(model_img, np.ndarray):
        result = model_img.copy()
    else:
        result = create_placeholder_image("Base Model")
    
    # Convert to PIL for drawing
    img_pil = Image.fromarray(result)
    draw = ImageDraw.Draw(img_pil)
    
    try:
        font = ImageFont.truetype("arial.ttf", 16)
        small_font = ImageFont.truetype("arial.ttf", 12)
    except:
        font = ImageFont.load_default()
        small_font = ImageFont.load_default()
    
    # Add overlay showing what's being applied
    overlay = Image.new('RGBA', img_pil.size, (0,0,0,0))
    overlay_draw = ImageDraw.Draw(overlay)
    overlay_draw.rectangle([10, 10, 390, 100], fill=(0,100,0,128))
    
    img_pil = Image.alpha_composite(img_pil.convert('RGBA'), overlay).convert('RGB')
    draw = ImageDraw.Draw(img_pil)
    
    y_pos = 20
    draw.text((20, y_pos), f"Base Model: {model_source}", fill='white', font=small_font)
    y_pos += 25
    
    if garment_top is not None:
        draw.text((20, y_pos), "✓ Top garment applied", fill='lightgreen', font=font)
        y_pos += 25
    if garment_bottom is not None:
        draw.text((20, y_pos), "✓ Bottom garment applied", fill='lightgreen', font=font)
    
    return np.array(img_pil)

def edit_outfit_with_text(current_image, edit_prompt, strength=0.7):
    """Edit the current outfit using text prompt"""
    if current_image is None:
        return None, "Please generate an outfit first!"
    
    if not edit_prompt.strip():
        return current_image, "Please enter an edit prompt!"
    
    try:
        # Convert to PIL if numpy array
        if isinstance(current_image, np.ndarray):
            image = Image.fromarray(current_image)
        else:
            image = current_image
            
        # Resize if needed
        image = image.resize((400, 600))
        
        # Create a simple mask for clothing area
        mask = Image.new('L', (400, 600), 0)
        mask_draw = ImageDraw.Draw(mask)
        mask_draw.rectangle([50, 150, 350, 450], fill=255)  # Clothing area
        
        if pipe_inpaint:
            prompt = f"person wearing {edit_prompt}, high quality fashion photo"
            negative_prompt = "blurry, low quality, distorted, nude"
            
            edited_image = pipe_inpaint(
                prompt=prompt,
                image=image,
                mask_image=mask,
                negative_prompt=negative_prompt,
                num_inference_steps=20,
                strength=strength,
                guidance_scale=7.5
            ).images[0]
            
            return np.array(edited_image), f"✅ Outfit edited: {edit_prompt}"
        else:
            # Fallback: add text overlay
            draw = ImageDraw.Draw(image)
            draw.text((10, 10), f"Edit: {edit_prompt}", fill='red')
            return np.array(image), f"✅ Edit applied: {edit_prompt}"
            
    except Exception as e:
        return current_image, f"❌ Error editing: {str(e)}"

# Initialize models
print("Loading AI models...")
models_loaded = initialize_models()
print(f"Models loaded: {models_loaded}")

# Create the main interface
with gr.Blocks(css=".output-image, .input-image, .image-preview {height: 400px !important}", 
               title="AI Outfit Designer") as demo:
    
    # Header
    gr.HTML("""
    <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
        <div>
            <h1>👗 AI Outfit Designer & Virtual Try-On 👔</h1>
            <h4>v1.0 - Upload your own photo or choose from our models</h4>
            <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
                <span style="background: #4CAF50; color: white; padding: 5px 10px; border-radius: 15px; margin: 5px;">
                    ✨ AI-Powered
                </span>
                <span style="background: #2196F3; color: white; padding: 5px 10px; border-radius: 15px; margin: 5px;">
                    🎨 Text Editing
                </span>
                <span style="background: #FF9800; color: white; padding: 5px 10px; border-radius: 15px; margin: 5px;">
                    👗 Virtual Try-On
                </span>
            </div>
        </div>
    </div>
    """)
    
    with gr.Row():
        # Left Column - Model Selection
        with gr.Column():
            gr.HTML("""
            <div style="text-align: center;">
                <h3>👤 Step 1: Choose Your Base Model</h3>
                <p>Upload your own photo or select a base model to dress up</p>
            </div>
            """)
            
            model_image = gr.Image(
                sources=['upload', 'clipboard'], 
                type="numpy", 
                label="Upload Your Photo (Optional)",
                height=400
            )
            
            # Model selection dropdown
            model_selector = gr.Dropdown(
                choices=list(MODEL_AVATARS.keys()),
                value="Fashion Model A",
                label="Or Choose a Base Model",
                interactive=True
            )
            
            # Example models gallery (shows actual base model images)
            gr.Markdown("### 📸 Available Base Models")
            
            # Create a function to show model preview
            def show_model_preview(model_name):
                if model_name in MODEL_AVATARS:
                    return load_model_image_from_url(MODEL_AVATARS[model_name])
                return create_placeholder_image("Select a model")
            
            model_preview = gr.Image(label="Selected Base Model Preview", height=300, interactive=False)
            
            # Update preview when dropdown changes
            model_selector.change(fn=show_model_preview, inputs=[model_selector], outputs=[model_preview])
        
        # Middle Column - Garments and Style
        with gr.Column():
            gr.HTML("""
            <div style="text-align: center;">
                <h3>👕 Step 2: Choose Style or Upload Garments</h3>
                <p>Upload specific garments or describe your desired style</p>
            </div>
            """)
            
            # Style prompt (new feature)
            style_prompt = gr.Textbox(
                label="Describe Your Desired Style",
                placeholder="e.g., 'casual jeans and t-shirt', 'elegant black dress', 'business suit'",
                lines=2
            )
            
            with gr.Row():
                garment_top = gr.Image(
                    sources='upload', 
                    type="numpy", 
                    label="Top Garment (Optional)",
                    height=200
                )
                garment_bottom = gr.Image(
                    sources='upload', 
                    type="numpy", 
                    label="Bottom Garment (Optional)",
                    height=200
                )
            
            # Quick style buttons
            gr.Markdown("### 🎯 Quick Style Options")
            with gr.Row():
                casual_btn = gr.Button("👕 Casual", size="sm")
                business_btn = gr.Button("💼 Business", size="sm")
                formal_btn = gr.Button("🎩 Formal", size="sm")
                trendy_btn = gr.Button("✨ Trendy", size="sm")
            
            generate_btn = gr.Button("🎨 Generate Outfit", variant="primary", size="lg")
        
        # Right Column - Results
        with gr.Column():
            gr.HTML("""
            <div style="text-align: center;">
                <h3>✨ Your Generated Outfit</h3>
            </div>
            """)
            
            result_image = gr.Image(label="Generated Outfit", height=400)
            
            generation_status = gr.Markdown("Upload a photo or select a model, then click Generate!")
    
    # Separator
    gr.HTML("<hr style='margin: 30px 0; border: 1px solid #ddd;'>")
    
    # Bottom Section - Text Editor
    gr.HTML("""
    <div style="text-align: center; margin: 20px;">
        <h2>✏️ Step 3: Edit Your Outfit with AI</h2>
        <p>Use natural language to modify your generated outfit</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column():
            edit_prompt = gr.Textbox(
                label="Describe Your Edits",
                placeholder="e.g., 'change shirt to red', 'add a jacket', 'make it more colorful'",
                lines=3
            )
            
            with gr.Row():
                edit_strength = gr.Slider(
                    minimum=0.3,
                    maximum=1.0,
                    value=0.7,
                    step=0.1,
                    label="Edit Strength"
                )
            
            edit_btn = gr.Button("🖌️ Apply Edit", variant="secondary", size="lg")
            
            # Quick edit options
            gr.Markdown("### ⚡ Quick Edits")
            with gr.Row():
                add_jacket_btn = gr.Button("Add Jacket", size="sm")
                change_color_btn = gr.Button("Change Colors", size="sm")
                make_formal_btn = gr.Button("Make Formal", size="sm")
                add_accessories_btn = gr.Button("Add Accessories", size="sm")
        
        with gr.Column():
            edited_result = gr.Image(label="Edited Outfit", height=400)
            edit_status = gr.Markdown("Generate an outfit first, then describe your edits!")
    
    # Examples Section
    gr.HTML("<hr style='margin: 30px 0; border: 1px solid #ddd;'>")
    gr.Markdown("## 📚 Examples & Tips")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            ### 🎯 Style Prompt Examples:
            - "Casual weekend outfit with jeans"
            - "Professional business attire" 
            - "Elegant evening dress for dinner"
            - "Sporty athleisure wear"
            - "Vintage 1950s inspired look"
            """)
        
        with gr.Column():
            gr.Markdown("""
            ### ✏️ Edit Prompt Examples:
            - "Change the shirt to a red color"
            - "Add a denim jacket over the outfit"
            - "Make the dress shorter and more casual"
            - "Add gold jewelry and accessories"
            - "Change to a winter coat and boots"
            """)
    
    # Store current image for editing
    current_generated = gr.State(None)
    
    # Event Handlers
    def generate_outfit(model_img, garment1, garment2, style_text, selected_model):
        result = get_outfit_tryon_result(model_img, garment1, garment2, style_text, selected_model)
        
        if model_img is not None:
            model_info = "uploaded photo"
        elif selected_model:
            model_info = selected_model
        else:
            model_info = "no model selected"
            
        status = f"✅ Generated outfit on {model_info}"
        if style_text:
            status += f" with style: {style_text}"
        if garment1 is not None or garment2 is not None:
            status += f" using uploaded garments"
            
        return result, result, status
    
    def apply_edit(current_img, edit_text, strength):
        if current_img is None:
            return None, "❌ Please generate an outfit first!"
        
        edited_img, status = edit_outfit_with_text(current_img, edit_text, strength)
        return edited_img, status
    
    # Connect events
    generate_btn.click(
        fn=generate_outfit,
        inputs=[model_image, garment_top, garment_bottom, style_prompt, model_selector],
        outputs=[result_image, current_generated, generation_status]
    )
    
    edit_btn.click(
        fn=apply_edit,
        inputs=[current_generated, edit_prompt, edit_strength],
        outputs=[edited_result, edit_status]
    )
    
    # Quick style buttons
    casual_btn.click(lambda: "casual jeans and t-shirt outfit", outputs=[style_prompt])
    business_btn.click(lambda: "professional business suit", outputs=[style_prompt])
    formal_btn.click(lambda: "elegant formal evening wear", outputs=[style_prompt])
    trendy_btn.click(lambda: "trendy fashionable modern outfit", outputs=[style_prompt])
    
    # Quick edit buttons
    add_jacket_btn.click(lambda: "add a stylish jacket", outputs=[edit_prompt])
    change_color_btn.click(lambda: "change colors to be more vibrant", outputs=[edit_prompt])
    make_formal_btn.click(lambda: "make the outfit more formal and elegant", outputs=[edit_prompt])
    add_accessories_btn.click(lambda: "add fashionable accessories and jewelry", outputs=[edit_prompt])

if __name__ == "__main__":
    demo.launch()