import gradio as gr
import numpy as np
import random
from diffusers import DiffusionPipeline
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"

if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

####

import gradio as gr
import replicate


def generate_image(model, lora_scale, guidance_scale, prompt_strength, num_steps, prompt):
    output = replicate.run(
        "dd-ds-ai/lora_test_01:70c669221124d8aaf0fc494f9553468bd069483a19e74b2753262008b1e8fbb2",
        input={
            "model": model,
            "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
        }
    )
    image_url = output[0] if output else None
    return image_url


# Gradio-Interface erstellen
def create_gradio_interface():
    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")

    # Erstelle ein Button-Interface für die Bildgenerierung
    generate_btn = gr.Button("Bild generieren")

    # Gradio Interface
    interface = gr.Interface(
        fn=generate_image,  # Die Funktion, die aufgerufen wird
        inputs=[lora_scale, guidance_scale, prompt_strength, num_steps, prompt],  # Eingaben
        outputs=gr.Image(label="Generated Image"),  # Ausgabe als Bild
    )

    # Binde den Button an die Bildgenerierung
    interface.launch(share=True)


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


# demo.queue().launch()