Update README.md

README.md

@@ -75,60 +75,34 @@ import torch
 from PIL import Image
 from unipicv2.pipeline_stable_diffusion_3_kontext import StableDiffusion3KontextPipeline
 from unipicv2.transformer_sd3_kontext import SD3Transformer2DKontextModel
-from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL, BitsAndBytesConfig
+from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL
 from transformers import CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
 
 # Load model components
-pretrained_model_name_or_path = "/mnt/datasets_vlm/chris/hf_ckpt/Unipic2-t2i"
-# int4 is recommended for inference:lower VRAM with no quality loss  {"int4", "fp16"} 
-quant = "int4"  
-
-# BitsAndBytes config
-bnb4 = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_use_double_quant=True,         
-    bnb_4bit_quant_type="nf4",              
-    bnb_4bit_compute_dtype=torch.bfloat16, 
-)
-bnb8 = BitsAndBytesConfig(load_in_8bit=True)
-
-if quant == "int4":
-    transformer = SD3Transformer2DKontextModel.from_pretrained(
-        pretrained_model_name_or_path, subfolder="transformer",
-        quantization_config=bnb4, device_map="auto", low_cpu_mem_usage=True
-    ).cuda()
-    text_qconf = bnb8
-    vae_dtype = torch.float16
-else:  # fp16
-    transformer = SD3Transformer2DKontextModel.from_pretrained(
-        pretrained_model_name_or_path, subfolder="transformer",
-        torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True
-    ).cuda()
-    text_qconf = None
-    vae_dtype = torch.float16
+pretrained_model_name_or_path = "Skywork/UniPic2-SD3.5M-Kontext-2B"
+
+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=vae_dtype, device_map="auto", low_cpu_mem_usage=True
-)
+    torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True
+).cuda()
 
 # Load text encoders
 text_encoder = CLIPTextModelWithProjection.from_pretrained(
-    pretrained_model_name_or_path, subfolder="text_encoder",
-    quantization_config=text_qconf, torch_dtype=None, device_map="auto", low_cpu_mem_usage=True
-)
+    pretrained_model_name_or_path, subfolder="text_encoder", torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True
+).cuda()
 tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
 
 text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
-    pretrained_model_name_or_path, subfolder="text_encoder_2",
-    quantization_config=text_qconf, torch_dtype=None, device_map="auto", low_cpu_mem_usage=True
-)
+    pretrained_model_name_or_path, subfolder="text_encoder_2", torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True
+).cuda()
 tokenizer_2 = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer_2")
 
 text_encoder_3 = T5EncoderModel.from_pretrained(
-    pretrained_model_name_or_path, subfolder="text_encoder_3",
-    quantization_config=text_qconf, torch_dtype=None, device_map="auto", low_cpu_mem_usage=True
-)
+    pretrained_model_name_or_path, subfolder="text_encoder_3", torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True
+).cuda()
 tokenizer_3 = T5TokenizerFast.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer_3")
 
 scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(