File size: 1,819 Bytes
3ac77f6 f182529 c2a9e07 f182529 3ac77f6 db8548c c2a9e07 db8548c 38c87cb db8548c c2a9e07 db8548c c2a9e07 db8548c c2a9e07 db8548c c2a9e07 164d77e db8548c c2a9e07 51e12d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from fix_cache import remove_old_cache
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
remove_old_cache()
# Load different models
models = {
"Stable Diffusion v1.5": "runwayml/stable-diffusion-v1-5",
"Stable Diffusion v2.1": "stabilityai/stable-diffusion-2-1",
"Anime Diffusion": "hakurei/waifu-diffusion-v1-4",
}
# Function to load the selected model
def load_model(model_name):
model_id = models[model_name]
pipe = StableDiffusionPipeline.from_pretrained(
model_id, torch_dtype=torch.float16
)
pipe = pipe.to("cpu") # Use GPU
return pipe
# Load the default model
current_pipe = load_model("Stable Diffusion v1.5")
# Function to generate image
def generate_image(prompt, model_name):
global current_pipe
# Reload pipeline if the model changes
if model_name not in current_pipe.config["_name_or_path"]:
current_pipe = load_model(model_name)
# Generate the image
image = current_pipe(prompt).images[0]
return image
# Create Gradio interface
with gr.Blocks() as demo:
gr.Markdown("### Multi-Model Text-to-Image Generator")
with gr.Row():
with gr.Column():
text_input = gr.Textbox(label="Enter a text prompt", placeholder="Describe the image you want...")
model_selector = gr.Dropdown(
label="Select Model", choices=list(models.keys()), value="Stable Diffusion v1.5"
)
generate_button = gr.Button("Generate Image")
with gr.Column():
output_image = gr.Image(label="Generated Image")
with gr.Column():
output_image2 = gr.Image(label= "Generated Image 2")
generate_button.click(
generate_image, inputs=[text_input, model_selector], outputs=output_image
)
demo.launch(share=True) |