fffiloni's picture
Create app.py
5494b47
raw
history blame
998 Bytes
import gradio as gr
import torch
from diffusers import DiffusionPipeline, AutoencoderKL
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipe = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
vae=vae, torch_dtype=torch.float16, variant="fp16",
use_safetensors=True
)
# This is where you load your trained weights
pipe.load_lora_weights("victor/outicon")
pipe.to("cuda")
def infer (prompt):
image = pipe(prompt=prompt, num_inference_steps=50).images[0]
return image
css = """
#col-container {max-width: 780px; margin-left: auto; margin-right: auto;}
"""
with gr.Blocks() as demo:
with gr.Column(elem_id="col-container"):
prompt_in = gr.Textbox(label="Prompt")
submit_btn = gr.Button("Submit")
image_out = gr.Image(label="Image output")
submit_btn.click(
fn = infer,
inputs = [prompt_in],
outputs = [image_out]
)
demo.queue().launch()