import torch
from diffusers import FluxPipeline
import gradio as gr

def generate_flux_image(
    prompt, 
    width=512, 
    height=512, 
    num_inference_steps=50,
    guidance_scale=7.5,
    seed=None
):
    # Zufallsgenerator-Initialisierung
    if seed is None:
        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
    else:
        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
    
    # Pipeline laden
    pipeline = FluxPipeline.from_pretrained(
        "improvements/flux",  # Passen Sie den Pfad an
        torch_dtype=torch.float16
    )
    
    # Überprüfen und ggf. auf GPU verschieben
    if torch.cuda.is_available():
        pipeline = pipeline.to("cuda")
    
    # Bildgenerierung
    image = pipeline(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator
    ).images[0]
    
    return image

# Gradio-Interface
def create_gradio_interface():
    with gr.Blocks() as demo:
        with gr.Row():
            prompt = gr.Textbox(label="Bildprompt")
            width = gr.Slider(minimum=64, maximum=2048, value=512, label="Breite")
            height = gr.Slider(minimum=64, maximum=2048, value=512, label="Höhe")
        
        with gr.Row():
            steps = gr.Slider(minimum=10, maximum=100, value=50, label="Inference Steps")
            guidance = gr.Slider(minimum=1, maximum=15, value=7.5, label="Guidance Scale")
            seed = gr.Number(label="Seed (optional)", precision=0)
        
        generate_btn = gr.Button("Bild generieren")
        output_image = gr.Image(label="Generiertes Bild")
        
        generate_btn.click(
            fn=generate_flux_image, 
            inputs=[prompt, width, height, steps, guidance, seed],
            outputs=output_image
        )
    
    return demo

# Interface starten
if __name__ == "__main__":
    interface = create_gradio_interface()
    interface.launch(share=True)