Create app.py

app.py

@@ -0,0 +1,512 @@
+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://images.unsplash.com/photo-1746527971885-001a9643e0a7?q=80&w=1074&auto=format&fit=crop&ixlib=rb-4.1.0&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D",
+    "Fashion Model B": "https://images.unsplash.com/photo-1656587132026-1d9cf800acc7?q=80&w=687&auto=format&fit=crop&ixlib=rb-4.1.0&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D",
+    "Fashion Model C": "https://plus.unsplash.com/premium_photo-1723478410866-67c032dba011?q=80&w=650&auto=format&fit=crop&ixlib=rb-4.1.0&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D",
+    "Casual Model A": "https://images.unsplash.com/photo-1737085725643-e751bcf7ff2b?q=80&w=627&auto=format&fit=crop&ixlib=rb-4.1.0&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D"
+}
+
+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()