Viewgen / app.py
RingL's picture
updated view app
7320387
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()