app.py

import spaces
import gradio as gr
from PIL import Image
import torch
import io
import base64
import tempfile
import os
from diffusers import AutoPipelineForInpainting, AutoencoderKL
import numpy as np

# Global variables for models
pipeline = None
segment_body = None

def load_models():
    """Load all required models"""
    global pipeline, segment_body
    
    print("🔄 Loading VAE...")
    vae = AutoencoderKL.from_pretrained(
        "madebyollin/sdxl-vae-fp16-fix",
        torch_dtype=torch.float16
    )
    
    print("🔄 Loading inpainting pipeline...")
    pipeline = AutoPipelineForInpainting.from_pretrained(
        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
        vae=vae,
        torch_dtype=torch.float16,
        variant="fp16",
        use_safetensors=True
    )
    
    pipeline = pipeline.to("cuda")
    
    print("🔄 Loading IP-Adapter...")
    pipeline.load_ip_adapter(
        "h94/IP-Adapter",
        subfolder="sdxl_models",
        weight_name="ip-adapter_sdxl.bin",
        low_cpu_mem_usage=True
    )
    
    print("🔄 Loading body segmentation...")
    try:
        from SegBody import segment_body as seg_func
        segment_body = seg_func
        print("✅ Body segmentation loaded!")
    except ImportError:
        print("⚠️ SegBody module not found, segmentation will be disabled")
    
    print("✅ All models loaded successfully!")

# Load models on startup
load_models()

def create_mask(person_img):
    """Generate body segmentation mask"""
    if segment_body is None:
        # Create a simple fallback mask (full body) if segmentation not available
        return Image.new('L', (512, 512), 255)
    
    try:
        # Try calling segment_body - it might expect a file path or PIL Image
        try:
            # First try with PIL Image directly
            seg_image, mask_img = segment_body(person_img, face=False)
        except (TypeError, AttributeError):
            # If that fails, try with file path
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
                temp_path = tmp_file.name
                person_img.save(temp_path)
            
            seg_image, mask_img = segment_body(temp_path, face=False)
            os.unlink(temp_path)
        
        # Ensure mask is PIL Image and resize
        if isinstance(mask_img, str):
            mask_img = Image.open(mask_img).convert('L')
        mask_img = mask_img.resize((512, 512))
        return mask_img
        
    except Exception as e:
        print(f"⚠️ Segmentation failed: {e}, using full mask")
        return Image.new('L', (512, 512), 255)

@spaces.GPU
def virtual_tryon(
    person_image,
    clothing_image,
    prompt="photorealistic, perfect body, beautiful skin, realistic skin, natural skin",
    negative_prompt="ugly, bad quality, bad anatomy, deformed body, deformed hands, deformed feet, deformed face, deformed clothing, deformed skin, bad skin, leggings, tights, stockings",
    ip_scale=0.8,
    strength=0.99,
    guidance_scale=7.5,
    num_steps=50
):
    """
    Virtual Try-On function
    
    Args:
        person_image: PIL Image of the person
        clothing_image: PIL Image of the clothing
        prompt: Generation prompt
        negative_prompt: Negative prompt
        ip_scale: IP-Adapter influence (0.0-1.0)
        strength: Inpainting strength (0.0-1.0)
        guidance_scale: CFG scale
        num_steps: Number of inference steps
    
    Returns:
        Generated PIL Image
    """
    try:
        if pipeline is None:
            raise ValueError("Models not loaded yet")
        
        print("📥 Processing images...")
        # Ensure images are PIL Images and resize
        if isinstance(person_image, np.ndarray):
            person_image = Image.fromarray(person_image)
        if isinstance(clothing_image, np.ndarray):
            clothing_image = Image.fromarray(clothing_image)
            
        person_img = person_image.convert('RGB').resize((512, 512))
        clothing_img = clothing_image.convert('RGB').resize((512, 512))
        
        # Generate body segmentation mask
        print("🎭 Generating segmentation mask...")
        mask_img = create_mask(person_img)
        
        # Set IP-Adapter scale
        pipeline.set_ip_adapter_scale(ip_scale)
        
        # Generate virtual try-on
        print("🎨 Generating virtual try-on...")
        result = pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=person_img,
            mask_image=mask_img,
            ip_adapter_image=clothing_img,
            strength=strength,
            guidance_scale=guidance_scale,
            num_inference_steps=num_steps,
        )
        
        generated_image = result.images[0]
        print("✅ Generation completed!")
        
        return generated_image
    
    except Exception as e:
        print(f"❌ Error: {str(e)}")
        raise gr.Error(f"Generation failed: {str(e)}")

# Create Gradio interface with simpler structure for better compatibility
demo = gr.Interface(
    fn=virtual_tryon,
    inputs=[
        gr.Image(label="Person Image", type="pil"),
        gr.Image(label="Clothing Image", type="pil"),
        gr.Textbox(
            label="Prompt",
            value="photorealistic, perfect body, beautiful skin, realistic skin, natural skin",
            lines=2
        ),
        gr.Textbox(
            label="Negative Prompt",
            value="ugly, bad quality, bad anatomy, deformed body, deformed hands, deformed feet, deformed face, deformed clothing, deformed skin, bad skin, leggings, tights, stockings",
            lines=2
        ),
        gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.8,
            step=0.05,
            label="IP-Adapter Scale (clothing influence)"
        ),
        gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.99,
            step=0.01,
            label="Inpainting Strength"
        ),
        gr.Slider(
            minimum=1.0,
            maximum=20.0,
            value=7.5,
            step=0.5,
            label="Guidance Scale"
        ),
        gr.Slider(
            minimum=20,
            maximum=100,
            value=50,
            step=5,
            label="Number of Steps"
        )
    ],
    outputs=gr.Image(label="Result", type="pil"),
    title="👕 Virtual Try-On",
    description="""
    Upload a photo of yourself and a clothing item to see how it looks on you!
    
    **Powered by Stable Diffusion XL + IP-Adapter + ZeroGPU**
    
    ### Tips for Best Results:
    - Use clear, well-lit photos
    - Person should be facing forward
    - Clothing image should show the item clearly
    - Higher steps = better quality but slower
    - Adjust IP-Adapter scale to control how much the clothing influences the result
    """,
    examples=None,
    cache_examples=False,
    allow_flagging="never"
)

# Launch the app
if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)