import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

def load_loras(pipe, lora_paths):
    if lora_paths.strip() == "":
        return pipe
    loras = [l.strip() for l in lora_paths.split(",")]
    return pipe

def generate_image(prompt, model_path, width, height, num_inference_steps, guidance_scale, lora_paths):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipe = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
    pipe = pipe.to(device)
    pipe = load_loras(pipe, lora_paths)
    image = pipe(prompt=prompt, width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
    return image

with gr.Blocks() as demo:
    gr.Markdown("## AI Image Generator")
    model_path = gr.Textbox(label="Model Path", value="path/to/model", lines=1)
    prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Enter your prompt...")
    width = gr.Slider(256, 1024, value=512, step=64, label="Width")
    height = gr.Slider(256, 1024, value=512, step=64, label="Height")
    guidance_scale = gr.Slider(1.0, 20.0, value=7.5, step=0.1, label="Guidance Scale")
    num_steps = gr.Slider(1, 150, value=50, step=1, label="Steps")
    lora_paths = gr.Textbox(label="LoRA Paths (comma separated)", value="", lines=1, placeholder="path1,path2,...")
    output_image = gr.Image(label="Generated Image")
    generate_button = gr.Button("Generate Image")
    generate_button.click(fn=generate_image, inputs=[prompt, model_path, width, height, num_steps, guidance_scale, lora_paths], outputs=output_image)
demo.launch()