import gradio as gr
from diffusers import StableDiffusionPipeline

def apply_lora(pipeline, lora_path):
    """
    Dummy function to simulate the application of LoRA weights.
    Replace this with your actual code to load and integrate LoRA weights.
    """
    if lora_path:
        print(f"Applying LoRA weights from {lora_path}")
        # Insert your LoRA integration code here.
    return pipeline

def generate_image(model_name, lora_path, width, height, inference_steps, prompt):
    # Use the provided model name or fall back to a default model
    model_id = model_name.strip() if model_name.strip() else "CompVis/stable-diffusion-v1-5"
    
    # Load the diffusion pipeline from Hugging Face
    pipeline = StableDiffusionPipeline.from_pretrained(model_id)
    device = "cuda" if gr.get_config().get("device") == "gpu" else "cpu"
    pipeline = pipeline.to(device)
    
    # Apply LoRA if a path is provided
    if lora_path.strip():
        pipeline = apply_lora(pipeline, lora_path.strip())
    
    # Generate the image using the specified parameters
    result = pipeline(prompt, width=width, height=height, num_inference_steps=inference_steps)
    return result.images[0]

# Build the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Image Generator with Custom Model & LoRA")
    
    model_name_box = gr.Textbox(
        label="Enter Model Name/ID (e.g., CompVis/stable-diffusion-v1-5)",
        value="CompVis/stable-diffusion-v1-5",
        lines=1
    )
    lora_path_box = gr.Textbox(
        label="Enter LoRA Path (leave empty if not using)",
        value="",
        lines=1
    )
    
    width_slider = gr.Slider(minimum=256, maximum=1024, value=512, step=64, label="Image Width")
    height_slider = gr.Slider(minimum=256, maximum=1024, value=512, step=64, label="Image Height")
    steps_slider = gr.Slider(minimum=10, maximum=100, value=50, step=1, label="Inference Steps")
    prompt_box = gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt")
    
    generate_button = gr.Button("Generate Image")
    output_image = gr.Image(label="Generated Image")
    
    generate_button.click(
        fn=generate_image, 
        inputs=[model_name_box, lora_path_box, width_slider, height_slider, steps_slider, prompt_box], 
        outputs=output_image
    )

demo.launch()