Spaces:
Sleeping
Sleeping
| 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() |