import gradio as gr
from diffusers import AutoPipelineForInpainting, AutoencoderKL
from diffusers.utils import load_image
import torch
from PIL import Image
import spaces
from SegBody import segment_body


# Load models

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float32)
pipeline = AutoPipelineForInpainting.from_pretrained(
    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", 
    vae=vae, 
    torch_dtype=torch.float32, 
    variant="fp16", 
    use_safetensors=True,
    device="cuda"
)
pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin", low_cpu_mem_usage=True)

# Function to process images
def virtual_try_on(img, clothing, prompt, negative_prompt, ip_scale=1.0, strength=0.99, guidance_scale=7.5, steps=100):
    _, mask_img = segment_body(img, face=False)
    pipeline.set_ip_adapter_scale(ip_scale)
    images = pipeline(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=img,
        mask_image=mask_img,
        ip_adapter_image=clothing,
        strength=strength,
        guidance_scale=guidance_scale,
        num_inference_steps=steps,
    ).images
    return images[0]

@spaces.GPU(duration=120)
def process_images(image, ip_image):
    image = image.convert("RGB").resize((512, 512))
    ip_image = ip_image.convert("RGB").resize((512, 512))
    
    seg_image, mask_image = segment_body(image, face=False)
    mask_image.resize((512, 512)) 
    

    return virtual_try_on(img=image,
               clothing=ip_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")

# Create the Gradio interface
interface = gr.Interface(
    fn=process_images, 
    inputs=[gr.Image(type="pil"), gr.Image(type="pil")], 
    outputs=gr.Image(type="pil"),
    title="Image Inpainting Demo",
    description="Upload two images for inpainting using Stable Diffusion XL."
)

# Launch the Gradio interface
interface.launch()