import gradio as gr
from diffusers import DiffusionPipeline
import torch

def generate_image(
    prompt, 
    width=512, 
    height=512, 
    num_inference_steps=50,
    guidance_scale=7.5
):
    try:
        # Alternativen Flux-Modell verwenden
        pipe = DiffusionPipeline.from_pretrained("enhanceaiteam/Flux-uncensored")
        
        # GPU-Optimierung
        if torch.cuda.is_available():
            pipe = pipe.to("cuda")
        
        image = pipe(
            prompt=prompt, 
            width=width, 
            height=height,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale
        ).images[0]
        return image
    except Exception as e:
        print(f"Fehler bei Bildgenerierung: {e}")
        return None

# Gradio-Interface
def create_gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# Flux Bildgenerator")
        
        with gr.Row():
            prompt = gr.Textbox(label="Bildprompt", lines=2)
            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")
        
        generate_btn = gr.Button("Bild generieren", variant="primary")
        output_image = gr.Image(label="Generiertes Bild")
        
        generate_btn.click(
            fn=generate_image, 
            inputs=[prompt, width, height, steps, guidance],
            outputs=output_image
        )
    
    return demo

# Hauptausführung
def main():
    interface = create_gradio_interface()
    interface.launch(share=True)

if __name__ == "__main__":
    main()