File size: 2,172 Bytes
392a2bc |
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 56 57 58 59 60 |
import os
import torch
import gradio as gr
import numpy as np
from PIL import Image
from diffusers import StableDiffusionPipeline,UNet2DConditionModel
NEGATIVE_PROMPT = "worst quality, low quality, bad anatomy, watermark, text, blurry, cartoon, unreal"
unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5",subfolder='unet').to("cuda")
# unet.load_lora_weights("./exp_output/celeba_finetune/checkpoint-20000", weight_name="pytorch_lora_weights.safetensors")
pipeline = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
unet=unet)
pipeline.load_lora_weights("./exp_output/celeba_finetune/checkpoint-20000", weight_name="pytorch_lora_weights.safetensors")
# Define a function to process input and return output
def generate_image(text,num_batch,is_use_lora,num_inference_steps):
# Process text to generate image
if is_use_lora:
pipeline.enable_lora()
else:
pipeline.disable_lora()
print('begin inference with text:', text, 'is_use_lora:', is_use_lora)
image = pipeline(text,
num_inference_steps=num_inference_steps,
num_images_per_prompt=num_batch,
negative_prompt=NEGATIVE_PROMPT).images
return image
with gr.Blocks() as demo:
with gr.Row():
with gr.Column():
with gr.Row():
is_use_lora = gr.Checkbox(label="Use LoRA", value=False)
num_batch = gr.Number(value=4,label="Number of batch")
num_inference_steps = gr.Number(value=20,label="Number of inference steps")
text_input = gr.Textbox(lines=2, label="Input text", value="A young woman with long hair and a big smile.")
generate_button = gr.Button(value="Generate image")
# image_out = gr.Image(label="Output image", height=512,width=512)
image_out = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", object_fit="contain", height="512")
generate_button.click(generate_image, inputs=[text_input,num_batch,is_use_lora,num_inference_steps], outputs=image_out)
demo.launch(server_port=7861)
|