Create generate_white.py

generate_white.py

@@ -0,0 +1,41 @@
+from diffusers import DiffusionPipeline
+import torch
+from PIL import Image
+import os
+
+def test_lora(lcm_speedup=False):
+    pipe = DiffusionPipeline.from_pretrained("/mnt/workspace/Project/VideoGen/ckpt/stable-diffusion-v1-5", torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False)
+    pipe.to("cuda")
+
+    lora_path = "/mnt/workspace/Project/VideoGen/TrainT2V/diffusers/examples/text_to_image/results" 
+    pipe.load_lora_weights(pretrained_model_name_or_path_or_dict=lora_path, weight_name="white_100epoch_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(["artifacts", "lcm"], adapter_weights=[1.0, 1.0])
+
+    prompts = [
+        "A central pure white artifact with a fluid, swirling form, resembling a twisted cloth. The artifact is highly detailed with intricate textures that catch the light, creating a sense of depth and realism. Set against a soft, out-of-focus black background that transitions from light to dark, the artifact stands out prominently."
+    ]
+
+    # 设置生成参数
+    if lcm_speedup:
+        num_inference_steps = 8
+        guidance_scale = 2
+    else:
+        num_inference_steps = 30
+        guidance_scale = 7.5
+        
+    num_samples_per_prompt = 10  # 生成10张图片
+
+    # 确保images文件夹存在
+    if not os.path.exists("images_result"):
+        os.makedirs("images_result")
+
+    # 为每个 prompt 生成 num_samples_per_prompt 张图片
+    for i, prompt in enumerate(prompts):
+        for j in range(num_samples_per_prompt):
+            image = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
+            image.save(f"images_result/{i}_{j}.png")  # 保存图片
+
+if __name__ == "__main__":
+    test_lora()