import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

def generate_image(prompt, guidance_scale=7.5, num_steps=30, lora_scale=0.8):
    try:
        # Initialize the pipeline
        pipe = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
        )
        
        # Load LoRA weights
        pipe.load_lora_weights("Scalino84/my-flux-face-v2")
        
        if torch.cuda.is_available():
            pipe = pipe.to("cuda")
        
        # Generate image
        image = pipe(
            prompt=prompt,
            num_inference_steps=int(num_steps),
            guidance_scale=float(guidance_scale),
        ).images[0]
        
        return image
    
    except Exception as e:
        print(f"Error: {str(e)}")
        raise gr.Error(str(e))

# Create Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", value="a photo of xyz person, professional headshot"),
        gr.Slider(minimum=1, maximum=20, value=7.5, label="Guidance Scale"),
        gr.Slider(minimum=20, maximum=100, value=30, label="Number of Steps"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.8, label="LoRA Scale")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Flux Face Generator",
    description="Generate images using your custom LoRA model"
)

# Launch the interface
interface.launch()