import gradio as gr
import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from PIL import Image
import numpy as np

def load_model():
    controlnet = ControlNetModel.from_pretrained("Kwai-Kolors/Kolors-Virtual-Try-On")
    pipe = StableDiffusionControlNetPipeline.from_pretrained(
        "Kwai-Kolors/Kolors-Virtual-Try-On",
        controlnet=controlnet,
        torch_dtype=torch.float16
    )
    if torch.cuda.is_available():
        pipe = pipe.to("cuda")
    return pipe

# Model'i global olarak yükle
try:
    model = load_model()
    print("Model başarıyla yüklendi!")
except Exception as e:
    print(f"Model yüklenirken hata: {str(e)}")

def virtual_try_on(person_image, garment_image):
    """
    Virtual try-on process
    """
    try:
        # Resimleri uygun formata dönüştür
        if person_image is None or garment_image is None:
            return None, "Error: Both images are required"
            
        # Model inference
        output = model(
            person_image,
            garment_image,
            num_inference_steps=30,
            guidance_scale=7.5
        )
        
        # Sonuç resmini al
        result_image = output.images[0]
        return result_image, "Success"
        
    except Exception as e:
        return None, f"Error: {str(e)}"

# Gradio arayüzü
demo = gr.Interface(
    fn=virtual_try_on,
    inputs=[
        gr.Image(type="pil", label="Person Image"),
        gr.Image(type="pil", label="Garment Image")
    ],
    outputs=[
        gr.Image(type="pil", label="Result"),
        gr.Text(label="Status")
    ],
    title="Virtual Try-On",
    description="Upload a person image and a garment image to see how the garment would look on the person."
)

if __name__ == "__main__":
    demo.launch()