from gradio_client import Client
import gradio as gr

client = Client("kratadata/Colabor-Image-Finetune")

def generate_image(
    prompt, lora_model, aspect_ratio,
    num_inference_steps, guidance_scale,
    seed, lora_scale
):
    result = client.predict(
        lora_model=lora_model,
        prompt=prompt,
        aspect_ratio=aspect_ratio,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        seed=seed,
        lora_scale=lora_scale,
        api_name="/generate"
    )
    return result

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", value="policewoman"),
        gr.Dropdown(["group1"], label="LoRA Model"),
        gr.Dropdown(["1:1", "16:9", "9:16"], label="Aspect Ratio", value="1:1"),
        gr.Slider(1, 100, step=1, label="Inference Steps", value=28),
        gr.Slider(0, 20, step=0.1, label="Guidance Scale", value=3),
        gr.Number(label="Seed (-1 for random)", value=-1),
        gr.Slider(0, 5, step=0.1, label="LoRA Scale", value=1),
    ],
    outputs=gr.Image(label="Generated Image"),
    title="LoRA Image Generator - kratadata/Colabor-Image-Finetune",
)

iface.launch()