from gradio_client import Client
import gradio as gr

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

def generate_image(prompt, aspect_ratio, lora_model):
    num_inference_steps = 28
    guidance_scale = 3
    seed = -1
    lora_scale = 1

    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

with gr.Blocks(css=".svelte-1ipelgc { display: none !important; }") as iface:
    gr.Markdown("# Gjeneruesi i Imazheve me Model të Personalizuar")
    
    prompt = gr.Textbox(label="Përshkrimi", placeholder="Shkruaj përshkrimin këtu...")
    aspect = gr.Dropdown(["1:1", "16:9", "9:16"], label="Dimensionet e fotos", value="1:1")
    model = gr.Dropdown(
        choices=["pizza-ananas", "pizza-pilz", "pizza-ananas-pilz"],
        label="Zgjidh Modelin e Imazhit"
    )
    btn = gr.Button("Gjenero")
    output = gr.Image(label="Rezultati")

    btn.click(fn=generate_image, inputs=[prompt, aspect, model], outputs=output)

iface.launch(share=True)