import gradio as gr
import replicate
import time


DEPLOYMENT_URIS = {
    "Lora 500": "dd-ds-ai/abendblatt-lora-500",
    "Lora 1000": "dd-ds-ai/abendblatt-lora-1000",
    "Lora 2000": "dd-ds-ai/abendblatt-lora-2000"
}


def generate_image(model_selection, lora_scale, guidance_scale, prompt_strength, num_steps, prompt):
    deployment_uri = DEPLOYMENT_URIS[model_selection]
    deployment = replicate.deployments.get(deployment_uri)

    prediction = deployment.predictions.create(
        input={
            "model": "dev",
            "lora_scale": lora_scale,
            "num_outputs": 1,
            "aspect_ratio": "1:1",
            "output_format": "webp",
            "guidance_scale": guidance_scale,
            "output_quality": 90,
            "prompt_strength": prompt_strength,
            "extra_lora_scale": 1,
            "num_inference_steps": num_steps,
            "prompt": prompt
        }
    )
    
    prediction.wait()
    output = prediction.output
    image_url = output[0] if output else None

    return image_url


# Gradio-Interface erstellen
def create_gradio_interface():
    model_selection = gr.Radio(choices=["Lora 500", "Lora 1000", "Lora 2000"], label="Model Selection", value="Lora 1000")
    
    lora_scale = gr.Slider(0, 2, value=1, step=0.1, label="Lora Scale")
    guidance_scale = gr.Slider(1, 10, value=3.5, step=0.1, label="Guidance Scale")
    prompt_strength = gr.Slider(0, 1, value=0.8, step=0.1, label="Prompt Strength")
    num_steps = gr.Slider(1, 50, value=28, step=1, label="Number of Inference Steps")
    prompt = gr.Textbox(label="Prompt", value="a person reading the hamburger abendblatt newspaper")

    generate_btn = gr.Button("Bild generieren")

    interface = gr.Interface(
        fn=generate_image,  
        inputs=[model_selection, lora_scale, guidance_scale, prompt_strength, num_steps, prompt], 
        outputs=gr.Image(label="Generated Image"),
    )

    interface.launch(share=True)


# Starte die Gradio-App
if __name__ == "__main__":
    create_gradio_interface()