Update README.md

README.md

@@ -81,7 +81,6 @@ from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL,BitsAndByte
 pretrained_model_name_or_path = "/path/to/UniPic2-Metaquery-Flash/UniPic2-Metaquery"
 vlm_path   = "/path/to/UniPic2-Metaquery-Flash/Qwen2.5-VL-7B-Instruct-AWQ"
 
-
 quant = "int4"  # {"int4", "fp16"}
 
 bnb4 = BitsAndBytesConfig(
@@ -93,12 +92,12 @@ bnb4 = BitsAndBytesConfig(
 
 if quant == "int4":
     transformer = SD3Transformer2DKontextModel.from_pretrained(
-        PRETRAINED_DIR, subfolder="transformer",
+        pretrained_model_name_or_path, subfolder="transformer",
         quantization_config=bnb4, device_map="auto", low_cpu_mem_usage=True
     )
 elif quant == "fp16":
     transformer = SD3Transformer2DKontextModel.from_pretrained(
-        PRETRAINED_DIR, subfolder="transformer",
+        pretrained_model_name_or_path, subfolder="transformer",
         torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True
     )
 else:
@@ -107,7 +106,7 @@ else:
 
 vae = AutoencoderKL.from_pretrained(
     pretrained_model_name_or_path, subfolder="vae",
-    torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True)
+    torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True).cuda()
 
 # Load Qwen2.5-VL model
 lmm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
@@ -120,8 +119,9 @@ processor.chat_template = processor.chat_template.replace(
     "{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}",
     "")
 
+# 加上cuda
 conditioner = StableDiffusion3Conditioner.from_pretrained(
-    pretrained_model_name_or_path, subfolder="conditioner",device_map="auto", torch_dtype=torch.bfloat16)
+    pretrained_model_name_or_path, subfolder="conditioner", torch_dtype=torch.float16).cuda()
 
 scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
 
@@ -167,6 +167,7 @@ image = pipeline(
 ).images[0]
 
 image.save("text2image.png")
+print(f"Image saved to text2image.png (quant={quant})")
 ```
 
 
@@ -189,7 +190,7 @@ min_pixels = max_pixels = int(image.height * 28 / 32 * image.width * 28 / 32)
 inputs = processor(
     text=texts, images=[image]*2,
     min_pixels=min_pixels, max_pixels=max_pixels,
-    videos=None, padding=True, return_tensors="pt")
+    videos=None, padding=True, return_tensors="pt").cuda()
 
 # Process with vision understanding
 input_ids, attention_mask, pixel_values, image_grid_thw = \
@@ -224,7 +225,9 @@ edited_image = pipeline(
     generator=torch.Generator(device=transformer.device).manual_seed(42)
 ).images[0]
 
-edited_image.save("image_editing.png")
+edited_image.save("edited_image.png")
+print(f"Image saved to edited_image.png (quant={quant})")
+
 ```