import gradio as gr
from gradio_client import Client, file

# Function to perform virtual try-on based on the selected model
def virtual_try_on(background, garment, garment_desc, denoise_steps, seed, model_choice):
    print(f"Model selected: {model_choice}")  # Debugging line to confirm the model choice
    
    # Initialize the client based on the selected model
    if model_choice == "IDM-VTON":
        client = Client("yisol/IDM-VTON")
    elif model_choice == "Virtual-Try-On":
        client = Client("Nymbo/Virtual-Try-On")
    else:
        raise ValueError("Model choice not recognized")  # Handle any unexpected model choice

    # Make the prediction using the selected model's API
    result = client.predict(
        dict={"background": file(background) if background else None, "layers": [], "composite": None},
        garm_img=file(garment),
        garment_des=garment_desc,
        is_checked=True,
        is_checked_crop=False,
        denoise_steps=denoise_steps,
        seed=seed,
        api_name="/tryon"
    )
    
    # Return the resulting images
    return result[0], result[1]  # First output: image, Second output: masked image

# Gradio interface
iface = gr.Interface(
    fn=virtual_try_on,
    inputs=[
        gr.Radio(choices=["IDM-VTON", "Virtual-Try-On"], label="اختر النموذج", value="IDM-VTON"),
        gr.Image(type="filepath", label="صورة الشخص"),
        gr.Image(type="filepath", label="صورة الملابس"),
        gr.Textbox(label="وصف الملابس (اختياري)"),
        gr.Slider(minimum=1, maximum=50, value=30, label="عدد خطوات التنقية"),
        gr.Slider(minimum=0, maximum=100, value=42, label="البذرة")
    ],
    outputs=[
        gr.Image(label="الصورة الناتجة"),
        gr.Image(label="الصورة المقنعة الناتجة")
    ],
    title="تجربة الملابس الافتراضية",
    description="اختر بين نموذجين لتجربة الملابس الافتراضية ورفع الصور."
)

# Launch the interface
iface.launch()