import os
import gradio as gr
from gradio_client import Client

api_key = os.getenv('HF_API_KEY')

custom_model = "lichorosario/dott_remastered_style_lora_sdxl"
weight_name = "dott_style.safetensors"

def predict(prompt):
    client = Client("fffiloni/sd-xl-custom-model", api_key=api_key)
    
    # Load the model
    client.predict(
        custom_model=custom_model,
        api_name="/load_model"
    )

    # Infer the image
    result = client.predict(
        custom_model=custom_model,
        weight_name=weight_name,
        prompt=prompt,
        inf_steps=25,
        guidance_scale=12,
        width=1024,
        height=512,
        seed=-1,
        lora_weight=1,
        api_name="/infer"
    )

    # Assuming result is a dictionary with keys 'image' and 'seed'
    image = result["image"]
    seed = result["seed"]

    return image, seed

with gr.Blocks() as demo:
    inputs = gr.Textbox(label="Prompt")
    outputs = [gr.Image(label="Image"), gr.Textbox(label="Seed")]
    greet_btn = gr.Button("Generate")
    greet_btn.click(fn=predict, inputs=inputs, outputs=outputs)

demo.launch()