Upload gradio_inference_t2i_lora.py

gradio_inference_t2i_lora.py

@@ -0,0 +1,59 @@
+
+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)
+