Update test_zimage.py

test_zimage.py

@@ -1,55 +1,55 @@
-import torch
-from diffusers import ZImagePipeline
-import os
-
-# 1. Load the pipeline
-# Use bfloat16 for optimal performance on supported GPUs
-pipe = ZImagePipeline.from_pretrained(
-    "../../../smodels/Z-Image-Turbo",
-    torch_dtype=torch.bfloat16,
-    low_cpu_mem_usage=False,
-)
-
-# =================【这里是新增的加载 LoRA 代码】=================
-# 指向你刚才训练输出的文件夹路径
-lora_dir = "./feifei-zimage-lora" 
-lora_file = "pytorch_lora_weights.safetensors"
-full_path = os.path.join(lora_dir, lora_file)
-
-if os.path.exists(full_path):
-    print(f"正在加载 LoRA: {full_path}")
-    try:
-        # adapter_name 可以随意起，用来标记这个 LoRA
-        pipe.load_lora_weights(lora_dir, weight_name=lora_file, adapter_name="feifei")
-        print("✅ LoRA 加载成功！")
-        
-        # [可选] 设置 LoRA 的权重强度 (1.0 = 100% 强度, 0.5 = 50%)
-        # pipe.set_adapters(["feifei"], adapter_weights=[1.0])
-        
-    except Exception as e:
-        print(f"❌ LoRA 加载失败: {e}")
-        print("可能是键名不匹配，或者文件损坏。")
-else:
-    print(f"⚠️ 找不到 LoRA 文件: {full_path}")
-# ===============================================================
-
-pipe.to("cuda")
-
-# [Optional] Attention Backend
-# pipe.transformer.set_attention_backend("flash")
-
-prompt = "jpop model in bikini at sea"
-
-# 2. Generate Image
-image = pipe(
-    prompt=prompt,
-    height=1024,
-    width=1024,
-    num_inference_steps=9,
-    guidance_scale=0.0,
-    generator=torch.Generator("cuda").manual_seed(42),
-    # cross_attention_kwargs={"scale": 1.0} # 另一种控制 LoRA 强度的方法
-).images[0]
-
-image.save("example_lora_test.png")
+import torch
+from diffusers import ZImagePipeline
+import os
+
+# 1. Load the pipeline
+# Use bfloat16 for optimal performance on supported GPUs
+pipe = ZImagePipeline.from_pretrained(
+    "./Z-Image-Turbo",
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=False,
+)
+
+# =================【这里是新增的加载 LoRA 代码】=================
+# 指向你刚才训练输出的文件夹路径
+lora_dir = "./feifei-zimage-lora" 
+lora_file = "pytorch_lora_weights.safetensors"
+full_path = os.path.join(lora_dir, lora_file)
+
+if os.path.exists(full_path):
+    print(f"正在加载 LoRA: {full_path}")
+    try:
+        # adapter_name 可以随意起，用来标记这个 LoRA
+        pipe.load_lora_weights(lora_dir, weight_name=lora_file, adapter_name="feifei")
+        print("✅ LoRA 加载成功！")
+        
+        # [可选] 设置 LoRA 的权重强度 (1.0 = 100% 强度, 0.5 = 50%)
+        # pipe.set_adapters(["feifei"], adapter_weights=[1.0])
+        
+    except Exception as e:
+        print(f"❌ LoRA 加载失败: {e}")
+        print("可能是键名不匹配，或者文件损坏。")
+else:
+    print(f"⚠️ 找不到 LoRA 文件: {full_path}")
+# ===============================================================
+
+pipe.to("cuda")
+
+# [Optional] Attention Backend
+# pipe.transformer.set_attention_backend("flash")
+
+prompt = "jpop model in bikini at sea"
+
+# 2. Generate Image
+image = pipe(
+    prompt=prompt,
+    height=1024,
+    width=1024,
+    num_inference_steps=9,
+    guidance_scale=0.0,
+    generator=torch.Generator("cuda").manual_seed(42),
+    # cross_attention_kwargs={"scale": 1.0} # 另一种控制 LoRA 强度的方法
+).images[0]
+
+image.save("example_lora_test.png")
 print("图像已保存为 example_lora_test.png")