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

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipeline = AutoPipelineForInpainting.from_pretrained(os.environ.get('MODEL'), vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
pipeline.load_ip_adapter(os.environ.get('IP_ADAPTER'), subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")

@spaces.GPU(enable_queue=True)
def squarify_image(img):
    if(img.height > img.width): bg_size = img.height
    else:  bg_size = img.width
    bg = Image.new(mode="RGB", size=(bg_size,bg_size), color="white")
    bg.paste(img, ( int((bg.width - bg.width)/2), 0) )

    return bg

@spaces.GPU(enable_queue=True)
def divisible_by_8(image):
    width, height = image.size
    
    # Calculate the new width and height that are divisible by 8
    new_width = (width // 8) * 8
    new_height = (height // 8) * 8
    
    # Resize the image
    resized_image = image.resize((new_width, new_height))
    
    return resized_image

@spaces.GPU(enable_queue=True)
def generate(person, clothing):
    person.thumbnail((1024,1024))
    person = divisible_by_8(person)
    
    clothing.thumbnail((1024,1024))
    clothing = divisible_by_8(clothing)

    image = squarify_image(person)

    seg_image, mask_image = segment_body(image, face=False)

    seg_cloth = segment_clothing(clothing, clothes= ["Upper-clothes", "Skirt", "Pants", "Dress", "Belt"])
    #seg_cloth = clothing
    
    pipeline.to("cuda")
    pipeline.set_ip_adapter_scale(1.0)
    images = pipeline(
        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",
        image=image,
        mask_image=mask_image,
        ip_adapter_image=seg_cloth,
        width=image.width,
        height=image.height,
        strength=0.99,
        guidance_scale=7.5,
        num_inference_steps=100,
    ).images
    
    final = images[0].crop((0, 0, person.width, person.height))
    
    return final

iface = gr.Interface(fn=generate, 
                     inputs=[gr.Image(label='Person', type='pil'), gr.Image(label='Clothing', type='pil')], 
                     outputs=[gr.Image(label='Result')],
                     title='Fashion Try-On',
                     description="""
                     by [Tony Assi](https://www.tonyassi.com/)

                     Check out [Virtual Try-On Pro](https://huggingface.co/spaces/tonyassi/Virtual-Try-On-Pro) !
                     
                     Please ❤️ this Space. I build custom AI apps for companies. <a href="mailto: tony.assi.media@gmail.com">Email me</a> for business inquiries.
                     """,
                     theme = gr.themes.Base(primary_hue="teal",secondary_hue="teal",neutral_hue="slate"),
                     examples=[["images/person1.jpg", "images/clothing1.jpg"], ["images/person1.jpg", "images/clothing2.jpg"]],)
iface.launch()