Update test_lora.py

test_lora.py

@@ -2,39 +2,55 @@ from diffusers import DiffusionPipeline
 import torch
 from PIL import Image
 
-# 加载 Stable Diffusion 模型
-pipe = DiffusionPipeline.from_pretrained("your_sd_dir/stable-diffusion-v1-5", torch_dtype=torch.float16)
-pipe.to("cuda")
-
-# 加载 LoRA 权重
-lora_path = "your_lora_dir"
-pipe.load_lora_weights(pretrained_model_name_or_path_or_dict=lora_path, weight_name="1epoch_lora.safetensors")
-
-# 定义 prompt 列表
-prompts = [
-    "Tang Dynasty Phoenix bird pattern, multi-integrated color complex figurative embroidery animal pattern, white background, asymmetry, meaning good weather, good luck, happy life. A symbol of good peace, abundance of children, supreme power and dominion. Worship of auspicious gods",
-    "Tang Dynasty Treasure Flower Pattern,flower,rotational,flower,rotational, radioactive arrangement,symmetry, solo, yellow theme"
-]
-
-# 设置生成参数
-num_inference_steps = 30
-guidance_scale = 7.5
-num_samples_per_prompt = 3
-
-# 创建一个空的图像列表
-all_images = []
-
-# 为每个 prompt 生成 num_samples_per_prompt 张图片
-for prompt in prompts:
-    images = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, num_images_per_prompt=num_samples_per_prompt).images
-    all_images.extend(images)
-
-# 创建一个 2x3 的网格图
-grid_image = Image.new('RGB', (3 * 512, 2 * 512))  # 假设每张图片大小为 512x512
-for idx, img in enumerate(all_images):
-    x = (idx % 3) * 512
-    y = (idx // 3) * 512
-    grid_image.paste(img, (x, y))
-
-# 保存网格图
-grid_image.save("test_lora_grid.png")
+
+def test_lora(lcm_speedup=Flase):
+    # 加载 Stable Diffusion 模型
+    pipe = DiffusionPipeline.from_pretrained("your_sd_dir/stable-diffusion-v1-5", torch_dtype=torch.float16, safety_checker = None, requires_safety_checker=False)
+    pipe.to("cuda")
+
+    # 加载 LoRA 权重
+    lora_path = "your_lora_dir" 
+    pipe.load_lora_weights(pretrained_model_name_or_path_or_dict=lora_path, weight_name="1epoch_lora.safetensors", adapter_name="pattern")
+    if lcm_speedup:
+        pipe.load_lora_weights(pretrained_model_name_or_path_or_dict=lora_path, weight_name="lcm_lora.safetensors", adapter_name="lcm")
+        pipe.set_adapters(["pattern", "lcm"], adapter_weights=[1.0, 1.0])
+
+    # 定义 prompt 列表
+    prompts = [
+        "Tang Dynasty Phoenix bird pattern, multi-integrated color complex figurative embroidery animal pattern, white background, asymmetry, meaning good weather, good luck, happy life. A symbol of good peace, abundance of children, supreme power and dominion. Worship of auspicious gods",
+        "Tang Dynasty Treasure Flower Pattern,flower,rotational,flower,rotational, radioactive arrangement,symmetry, solo, yellow theme"
+    ]
+
+    # 设置生成参数
+    if lcm_speedup:
+        num_inference_steps = 8
+        guidance_scale = 2
+    else:
+        num_inference_steps = 30
+        guidance_scale = 7.5
+        
+    num_samples_per_prompt = 3
+
+    # 创建一个空的图像列表
+    all_images = []
+
+    # 为每个 prompt 生成 num_samples_per_prompt 张图片
+    for prompt in prompts:
+        images = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, num_images_per_prompt=num_samples_per_prompt).images
+        all_images.extend(images)
+
+    # 创建一个 2x3 的网格图
+    grid_image = Image.new('RGB', (3 * 512, 2 * 512))  # 假设每张图片大小为 512x512
+    for idx, img in enumerate(all_images):
+        x = (idx % 3) * 512
+        y = (idx // 3) * 512
+        grid_image.paste(img, (x, y))
+
+    # 保存网格图
+    n = 4 if lcm_speedup else 30
+    grid_image.save(f"test_lora_grid_{n}_steps.png")
+
+
+if __name__ == "__main__":
+    test_lora()
+    test_lora(lcm_speedup=True)