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("""
Upload your own photo or select a base model to dress up
Upload specific garments or describe your desired style
Use natural language to modify your generated outfit