Spaces:
Sleeping
Sleeping
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() |