Update README.md

README.md

@@ -78,27 +78,50 @@ from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor
 from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL
 
 # Load model components
-pretrained_model_name_or_path = "/path/to/unipicv2_qwen2_5_vl_7b_sd_3_5m_kontext"
+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(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,  # 与 LMM/Cond 对齐
+)
+
+if quant == "int4":
+    transformer = SD3Transformer2DKontextModel.from_pretrained(
+        PRETRAINED_DIR, subfolder="transformer",
+        quantization_config=bnb4, device_map="auto", low_cpu_mem_usage=True
+    )
+elif quant == "fp16":
+    transformer = SD3Transformer2DKontextModel.from_pretrained(
+        PRETRAINED_DIR, subfolder="transformer",
+        torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True
+    )
+else:
+    raise ValueError(f"Unsupported quant: {quant}")
 
-transformer = SD3Transformer2DKontextModel.from_pretrained(
-    pretrained_model_name_or_path, subfolder="transformer", torch_dtype=torch.bfloat16).cuda()
 
 vae = AutoencoderKL.from_pretrained(
-    pretrained_model_name_or_path, subfolder="vae", torch_dtype=torch.bfloat16).cuda()
+    pretrained_model_name_or_path, subfolder="vae",
+    torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True)
 
 # Load Qwen2.5-VL model
 lmm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    "Qwen/Qwen2.5-VL-7B-Instruct",
-    torch_dtype=torch.bfloat16,
-    attn_implementation="flash_attention_2").cuda()
+    vlm_path,
+    torch_dtype=torch.bfloat16,device_map="auto",
+    attn_implementation="flash_attention_2")
 
-processor = Qwen2_5_VLProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
+processor = Qwen2_5_VLProcessor.from_pretrained(vlm_path)
 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 %}",
     "")
 
 conditioner = StableDiffusion3Conditioner.from_pretrained(
-    pretrained_model_name_or_path, subfolder="conditioner", torch_dtype=torch.bfloat16).cuda()
+    pretrained_model_name_or_path, subfolder="conditioner",device_map="auto", torch_dtype=torch.bfloat16)
 
 scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
 
@@ -166,7 +189,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").to("cuda")
+    videos=None, padding=True, return_tensors="pt")
 
 # Process with vision understanding
 input_ids, attention_mask, pixel_values, image_grid_thw = \