Views
No views yet
1CACHE_ROOT = pathlib.Path("qwen-image-int8-quanto ") # where we store INT8 modules
2TRANSFORMER_DIR = CACHE_ROOT / "qwen_image_transformer_int8"
3TEXT_ENCODER_DIR = CACHE_ROOT / "qwen_text_encoder_int8"
4
5def load_quantized_modules(transformer_dir: pathlib.Path, text_encoder_dir: pathlib.Path):
6 """
7 Load quantized modules (we saved them with the exact filenames the loaders expect).
8 """
9 tr = torch.load(str(transformer_dir / 'pytorch_model.bin'), weights_only=False)
10 te = torch.load(str(text_encoder_dir / 'pytorch_model.bin'), weights_only=False)
11 return tr, te
12
13def build_pipe(cls, transformer_dir: pathlib.Path, text_encoder_dir: pathlib.Path):
14 """
15 Build a pipeline of class `cls` by loading the quantized modules from disk.
16 Fresh module instances each time avoids offload-hook/state reuse hangs.
17 """
18 transformer, text_encoder = load_quantized_modules(transformer_dir, text_encoder_dir)
19 pipe = cls.from_pretrained(
20 BASE_MODEL_ID,
21 transformer=transformer,
22 text_encoder=text_encoder,
23 torch_dtype=torch.bfloat16,
24 use_safetensors=True,
25 low_cpu_mem_usage=True,
26 )
27 pipe.enable_model_cpu_offload()
28 pipe.set_progress_bar_config(disable=False)
29 return pipe