import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
from huggingface_hub import HfFolder

def generate_image(prompt, guidance_scale, num_steps, lora_scale):
    # Lade das Base Model
    pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float16
    ).to("cuda")
    
    # Lade dein LoRA
    pipe.load_lora_weights(
        "Scalino84/my-flux-face-v2",
        weight_name="flux_train_replicate.safetensors"
    )
    
    # Generiere das Bild
    image = pipe(
        prompt=prompt,
        num_inference_steps=num_steps,
        guidance_scale=guidance_scale,
        cross_attention_kwargs={"scale": lora_scale}
    ).images[0]
    
    return image

# Erstelle das Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Flux Face Generator")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                value="a photo of xyz person, professional headshot",
                lines=3
            )
            guidance = gr.Slider(
                label="Guidance Scale",
                minimum=1,
                maximum=20,
                value=7.5,
                step=0.5
            )
            steps = gr.Slider(
                label="Inference Steps",
                minimum=20,
                maximum=100,
                value=30,
                step=1
            )
            lora_scale = gr.Slider(
                label="LoRA Scale",
                minimum=0.1,
                maximum=1.0,
                value=0.8,
                step=0.1
            )
            generate = gr.Button("Generate Image")
        
        with gr.Column():
            output = gr.Image(label="Generated Image")
    
    generate.click(
        fn=generate_image,
        inputs=[prompt, guidance, steps, lora_scale],
        outputs=output
    )

demo.launch()