import gradio as gr
from diffusers import AutoPipelineForText2Image
import torch

# Load the base pipeline
base_model = "runwayml/stable-diffusion-v1-5"

# Define a function to create a pipeline with the appropriate LoRA weights
def create_pipeline_with_lora():
    pipeline = AutoPipelineForText2Image.from_pretrained(base_model, torch_dtype=torch.float32)
    pipeline.load_lora_weights("Front", weight_name="pytorch_lora_weights.safetensors")
    pipeline.load_lora_weights("Back", weight_name="pytorch_lora_weights.safetensors")
    pipeline.load_lora_weights("Rear", weight_name="pytorch_lora_weights.safetensors")
    return pipeline

pipeline = create_pipeline_with_lora()

# Define the prompt function
def generate_images(prompt):
    views = ["front", "back", "rear"]
    images = []
    for view in views:
        view_prompt = f"<a {view} view> of {prompt}"
        image = pipeline(view_prompt).images[0]
        images.append(image)
    return images

# Define the Gradio interface
def gradio_app(prompt_choice, custom_prompt):
    if prompt_choice == "Custom":
        prompt = custom_prompt
    else:
        prompt = prompt_choice
    return generate_images(prompt)

example_prompts = [
    "nighttime cityscape with a blue luxury car and a yellow sports car driving on a brightly lit street",
    "a white car and a gray car drive along a busy street"
]

with gr.Blocks() as demo:
    gr.Markdown("# Image View Generation with Diffusers")
    prompt_choice = gr.Dropdown(label="Select a Prompt", choices=example_prompts + ["Custom"], value="Custom")
    custom_prompt = gr.Textbox(label="Enter Custom Prompt", placeholder="Enter your own prompt here...")
    generate_button = gr.Button("Generate Images")
    output_gallery = gr.Gallery(label="Generated Images")

    generate_button.click(fn=gradio_app, inputs=[prompt_choice, custom_prompt], outputs=output_gallery)

# Launch the Gradio app
demo.launch()