Update app.py

app.py

@@ -1,512 +1,120 @@
-import gradio as gr
-import numpy as np
+import os
 import cv2
-from PIL import Image, ImageDraw, ImageFont
-import requests
+import numpy as np
 import json
 import random
+from PIL import Image, ImageDraw, ImageFont
+import asyncio
+
+import requests
 import base64
-import io
-import os
-from diffusers import StableDiffusionInpaintPipeline, StableDiffusionPipeline
-import torch
+import gradio as gr
 
-# Initialize AI models
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe_inpaint = None
-pipe_txt2img = None
+machine_number = 0
+model = os.path.join(os.path.dirname(__file__), "models/female_model.png")
 
-# 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"
+MODEL_MAP = {
+    "AI Model female_model": 'models/female_model.png',
+    "AI Model female_model1": 'models/female_model1.png',
+    "AI Model male_model": 'models/male_model.png',
+    "AI Model male_model2": 'models/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)
+def add_waterprint(img):
+    h, w, _ = img.shape
+    img = cv2.putText(img, 'Powered by OutfitAnyone', (int(0.3*w), h-20), cv2.FONT_HERSHEY_PLAIN, 2, (128, 128, 128), 2, 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_tryon_result(model_name, garment1, garment2, seed=1234):
+    # model_name = "AI Model " + model_name.split("\\")[-1].split(".")[0] # windows
+    model_name = "AI Model " + model_name.split("/")[-1].split(".")[0] # linux
+    print(model_name)
 
-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")
+    encoded_garment1 = cv2.imencode('.jpg', garment1)[1].tobytes()
+    encoded_garment1 = base64.b64encode(encoded_garment1).decode('utf-8')
 
-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()
+    if garment2 is not None:
+        encoded_garment2 = cv2.imencode('.jpg', garment2)[1].tobytes()
+        encoded_garment2 = base64.b64encode(encoded_garment2).decode('utf-8')
     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)
+        encoded_garment2 = ''
 
-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()
+    url = os.environ['OA_IP_ADDRESS']
+    headers = {'Content-Type': 'application/json'}
+    seed = random.randint(0, 1222222222)
+    data = {
+        "garment1": encoded_garment1,
+        "garment2": encoded_garment2,
+        "model_name": model_name,
+        "seed": seed
+    }
+    response = requests.post(url, headers=headers, data=json.dumps(data))
+    print("response code", response.status_code)
+    if response.status_code == 200:
+        result = response.json()
+        result = base64.b64decode(result['images'][0])
+        result_np = np.frombuffer(result, np.uint8)
+        result_img = cv2.imdecode(result_np, cv2.IMREAD_UNCHANGED)
     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)
+        print('server error!')
 
-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)}"
+    final_img = add_waterprint(result_img)
+    return final_img
 
-# 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;">
+with gr.Blocks(css = ".output-image, .input-image, .image-preview {height: 400px !important} ") as demo:
+    gr.HTML(
+        """
+        <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+        <a href="https://github.com/HumanAIGC/OutfitAnyone" style="margin-right: 20px; text-decoration: none; display: flex; align-items: center;">
+        </a>
         <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>
+            <h1 >Outfit Anyone: Ultra-high quality virtual try-on for Any Clothing and Any Person</h1>
+            <h4 >v0.9</h4>
+            <h5 style="margin: 0;">OutfitAnyone plus version is now online with any model and any cloth: https://www.outfitanyone.life/</h5>
+            <div style="display: flex; justify-content: center; align-items: center; text-align: center;>
+                <a href="https://arxiv.org/abs/2407.16224"><img src="https://img.shields.io/badge/Arxiv-2407.16224-red"></a>
+                <a href='https://humanaigc.github.io/outfit-anyone/'><img src='https://img.shields.io/badge/Project_Page-OutfitAnyone-green' alt='Project Page'></a>
+                <a href='https://github.com/HumanAIGC/OutfitAnyone'><img src='https://img.shields.io/badge/Github-Repo-blue'></a>
             </div>
         </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)
+            init_image = gr.Image(sources='clipboard', type="filepath", label="model", value=model)
             
-            # 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
-            )
-            
+            gr.HTML(
+                """
+                <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+                <div>
+                    <h3>Models are fixed and cannot be uploaded or modified; we only support users uploading their own garments.</h3>
+                    <h4 style="margin: 0;">For a one-piece dress or coat, you only need to upload the image to the 'top garment' section and leave the 'lower garment' section empty.</h4>
+                </div>
+                </div>
+                """)
             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")
+                garment_top = gr.Image(sources='upload', type="numpy", label="top garment")
+                garment_down = gr.Image(sources='upload', type="numpy", label="lower garment")
+
+            run_button = gr.Button(value="Run")
             
-            # 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])
+            gallery = gr.Image()
+
+            run_button.click(fn=get_tryon_result, 
+                             inputs=[
+                                    init_image,
+                                    garment_top,
+                                    garment_down,
+                                    ], 
+                             outputs=[gallery],
+                             concurrency_limit=2)
 
 if __name__ == "__main__":
+    ip = requests.get('http://ifconfig.me/ip', timeout=1).text.strip()
+    print("ip address alibaba", ip)
+    demo.queue(max_size=10)
     demo.launch()