Views
No views yet
python -m deepcompressor.app.diffusion.dataset.collect.calib svdq/flux.1-kontext-dev.yaml examples/diffusion/configs/collect/qdiff.yaml --pipeline-path svdq/flux.1-kontext-dev/In total 32 samples
Evaluating with batch size 1
Data: 3%|██▎ | 1/32 [13:57<7:12:32, 837.19s/it]
Sampling: 12%|█████████▍ | 1/8 [01:34<11:01, 94.44s/it]--save-model true or --save-model /PATH/TO/CHECKPOINT/DIRpython -m deepcompressor.app.diffusion.ptq svdq/flux.1-kontext-dev.yaml examples/diffusion/configs/svdquant/nvfp4.yaml --pipeline-path svdq/flux.1-kontext-dev/ --save-model ~/svdq/python -m deepcompressor.backend.nunchaku.convert --quant-path ~/svdq/ --output-root ~/svdq/ --model-name flux.1-kontext-dev-svdq-fp41 @staticmethod
2 def _default_build(
3 name: str, path: str, dtype: str | torch.dtype, device: str | torch.device, shift_activations: bool
4 ) -> DiffusionPipeline:
5 if not path:
6 if name == "sdxl":
7 path = "stabilityai/stable-diffusion-xl-base-1.0"
8 elif name == "sdxl-turbo":
9 path = "stabilityai/sdxl-turbo"
10 elif name == "pixart-sigma":
11 path = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"
12 elif name == "flux.1-kontext-dev":
13 path = "black-forest-labs/FLUX.1-Kontext-dev"
14 elif name == "flux.1-dev":
15 path = "black-forest-labs/FLUX.1-dev"
16 elif name == "flux.1-canny-dev":
17 path = "black-forest-labs/FLUX.1-Canny-dev"
18 elif name == "flux.1-depth-dev":
19 path = "black-forest-labs/FLUX.1-Depth-dev"
20 elif name == "flux.1-fill-dev":
21 path = "black-forest-labs/FLUX.1-Fill-dev"
22 elif name == "flux.1-schnell":
23 path = "black-forest-labs/FLUX.1-schnell"
24 else:
25 raise ValueError(f"Path for {name} is not specified.")
26 if name in ["flux.1-kontext-dev"]:
27 pipeline = FluxKontextPipeline.from_pretrained(path, torch_dtype=dtype)
28 elif name in ["flux.1-canny-dev", "flux.1-depth-dev"]:
29 pipeline = FluxControlPipeline.from_pretrained(path, torch_dtype=dtype)
30 elif name == "flux.1-fill-dev":
31 pipeline = FluxFillPipeline.from_pretrained(path, torch_dtype=dtype)
32 elif name.startswith("sana-"):
33 if dtype == torch.bfloat16:
34 pipeline = SanaPipeline.from_pretrained(path, variant="bf16", torch_dtype=dtype, use_safetensors=True)
35 pipeline.vae.to(dtype)
36 pipeline.text_encoder.to(dtype)
37 else:
38 pipeline = SanaPipeline.from_pretrained(path, torch_dtype=dtype)
39 else:
40 pipeline = AutoPipelineForText2Image.from_pretrained(path, torch_dtype=dtype)
41
42 # Debug output
43 print(">>> DEVICE:", device)
44 print(">>> PIPELINE TYPE:", type(pipeline))
45
46 # Try to move each component using .to_empty()
47 for name in ["unet", "transformer", "vae", "text_encoder"]:
48 module = getattr(pipeline, name, None)
49 if isinstance(module, torch.nn.Module):
50 try:
51 print(f">>> Moving {name} to {device} using to_empty()")
52 module.to_empty(device=device)
53 except Exception as e:
54 print(f">>> WARNING: {name}.to_empty({device}) failed: {e}")
55 try:
56 print(f">>> Falling back to {name}.to({device})")
57 module.to(device)
58 except Exception as ee:
59 print(f">>> ERROR: {name}.to({device}) also failed: {ee}")
60
61 # Identify main model (for patching)
62 model = getattr(pipeline, "unet", None) or getattr(pipeline, "transformer", None)
63 if model is not None:
64 replace_fused_linear_with_concat_linear(model)
65 replace_up_block_conv_with_concat_conv(model)
66 if shift_activations:
67 shift_input_activations(model)
68 else:
69 print(">>> WARNING: No model (unet/transformer) found for patching")
70
71 return pipeline1 @staticmethod
2 def _default_construct(
3 module: Attention,
4 /,
5 parent: tp.Optional["DiffusionTransformerBlockStruct"] = None,
6 fname: str = "",
7 rname: str = "",
8 rkey: str = "",
9 idx: int = 0,
10 **kwargs,
11 ) -> "DiffusionAttentionStruct":
12 if isinstance(module, FluxAttention):
13 # FluxAttention has different attribute names than standard attention
14 with_rope = True
15 num_query_heads = module.heads # FluxAttention uses 'heads', not 'num_heads'
16 num_key_value_heads = module.heads # FLUX typically uses same for q/k/v
17
18 # FluxAttention doesn't have 'to_out', but may have other output projections
19 # Check what output projection attributes actually exist
20 o_proj = None
21 o_proj_rname = ""
22
23 # Try to find the correct output projection
24 if hasattr(module, 'to_out') and module.to_out is not None:
25 o_proj = module.to_out[0] if isinstance(module.to_out, (list, tuple)) else module.to_out
26 o_proj_rname = "to_out.0" if isinstance(module.to_out, (list, tuple)) else "to_out"
27 elif hasattr(module, 'to_add_out'):
28 o_proj = module.to_add_out
29 o_proj_rname = "to_add_out"
30
31 q_proj, k_proj, v_proj = module.to_q, module.to_k, module.to_v
32 q_proj_rname, k_proj_rname, v_proj_rname = "to_q", "to_k", "to_v"
33 q, k, v = module.to_q, module.to_k, module.to_v
34 q_rname, k_rname, v_rname = "to_q", "to_k", "to_v"
35
36 # Handle the add_* projections that FluxAttention has
37 add_q_proj = getattr(module, "add_q_proj", None)
38 add_k_proj = getattr(module, "add_k_proj", None)
39 add_v_proj = getattr(module, "add_v_proj", None)
40 add_o_proj = getattr(module, "to_add_out", None)
41 add_q_proj_rname = "add_q_proj" if add_q_proj else ""
42 add_k_proj_rname = "add_k_proj" if add_k_proj else ""
43 add_v_proj_rname = "add_v_proj" if add_v_proj else ""
44 add_o_proj_rname = "to_add_out" if add_o_proj else ""
45
46 kwargs = (
47 "encoder_hidden_states",
48 "attention_mask",
49 "image_rotary_emb",
50 )
51 cross_attention = add_k_proj is not None
52 elif module.is_cross_attention:
53 q_proj, k_proj, v_proj = module.to_q, None, None
54 add_q_proj, add_k_proj, add_v_proj, add_o_proj = None, module.to_k, module.to_v, None
55 q_proj_rname, k_proj_rname, v_proj_rname = "to_q", "", ""
56 add_q_proj_rname, add_k_proj_rname, add_v_proj_rname, add_o_proj_rname = "", "to_k", "to_v", ""
57 else:
58 q_proj, k_proj, v_proj = module.to_q, module.to_k, module.to_v
59 add_q_proj = getattr(module, "add_q_proj", None)
60 add_k_proj = getattr(module, "add_k_proj", None)
61 add_v_proj = getattr(module, "add_v_proj", None)
62 add_o_proj = getattr(module, "to_add_out", None)
63 q_proj_rname, k_proj_rname, v_proj_rname = "to_q", "to_k", "to_v"
64 add_q_proj_rname, add_k_proj_rname, add_v_proj_rname = "add_q_proj", "add_k_proj", "add_v_proj"
65 add_o_proj_rname = "to_add_out"
66 if getattr(module, "to_out", None) is not None:
67 o_proj = module.to_out[0]
68 o_proj_rname = "to_out.0"
69 assert isinstance(o_proj, nn.Linear)
70 elif parent is not None:
71 assert isinstance(parent.module, FluxSingleTransformerBlock)
72 assert isinstance(parent.module.proj_out, ConcatLinear)
73 assert len(parent.module.proj_out.linears) == 2
74 o_proj = parent.module.proj_out.linears[0]
75 o_proj_rname = ".proj_out.linears.0"
76 else:
77 raise RuntimeError("Cannot find the output projection.")
78 if isinstance(module.processor, DiffusionAttentionProcessor):
79 with_rope = module.processor.rope is not None
80 elif module.processor.__class__.__name__.startswith("Flux"):
81 with_rope = True
82 else:
83 with_rope = False # TODO: fix for other processors
84 config = AttentionConfigStruct(
85 hidden_size=q_proj.weight.shape[1],
86 add_hidden_size=add_k_proj.weight.shape[1] if add_k_proj is not None else 0,
87 inner_size=q_proj.weight.shape[0],
88 num_query_heads=module.heads,
89 num_key_value_heads=module.to_k.weight.shape[0] // (module.to_q.weight.shape[0] // module.heads),
90 with_qk_norm=module.norm_q is not None,
91 with_rope=with_rope,
92 linear_attn=isinstance(module.processor, SanaLinearAttnProcessor2_0),
93 )
94 return DiffusionAttentionStruct(
95 module=module,
96 parent=parent,
97 fname=fname,
98 idx=idx,
99 rname=rname,
100 rkey=rkey,
101 config=config,
102 q_proj=q_proj,
103 k_proj=k_proj,
104 v_proj=v_proj,
105 o_proj=o_proj,
106 add_q_proj=add_q_proj,
107 add_k_proj=add_k_proj,
108 add_v_proj=add_v_proj,
109 add_o_proj=add_o_proj,
110 q=None, # TODO: add q, k, v
111 k=None,
112 v=None,
113 q_proj_rname=q_proj_rname,
114 k_proj_rname=k_proj_rname,
115 v_proj_rname=v_proj_rname,
116 o_proj_rname=o_proj_rname,
117 add_q_proj_rname=add_q_proj_rname,
118 add_k_proj_rname=add_k_proj_rname,
119 add_v_proj_rname=add_v_proj_rname,
120 add_o_proj_rname=add_o_proj_rname,
121 q_rname="",
122 k_rname="",
123 v_rname="",
124 )prompt or prompt_embeds. Cannot leave both prompt and prompt_embeds undefined.1def collect(config: DiffusionPtqRunConfig, dataset: datasets.Dataset):
2 samples_dirpath = os.path.join(config.output.root, "samples")
3 caches_dirpath = os.path.join(config.output.root, "caches")
4 os.makedirs(samples_dirpath, exist_ok=True)
5 os.makedirs(caches_dirpath, exist_ok=True)
6 caches = []
7
8 pipeline = config.pipeline.build()
9 model = pipeline.unet if hasattr(pipeline, "unet") else pipeline.transformer
10 assert isinstance(model, nn.Module)
11 model.register_forward_hook(CollectHook(caches=caches), with_kwargs=True)
12
13 batch_size = config.eval.batch_size
14 print(f"In total {len(dataset)} samples")
15 print(f"Evaluating with batch size {batch_size}")
16 pipeline.set_progress_bar_config(desc="Sampling", leave=False, dynamic_ncols=True, position=1)
17 for batch in tqdm(
18 dataset.iter(batch_size=batch_size, drop_last_batch=False),
19 desc="Data",
20 leave=False,
21 dynamic_ncols=True,
22 total=(len(dataset) + batch_size - 1) // batch_size,
23 ):
24 filenames = batch["filename"]
25 prompts = batch["prompt"]
26 seeds = [hash_str_to_int(name) for name in filenames]
27 generators = [torch.Generator(device=pipeline.device).manual_seed(seed) for seed in seeds]
28 pipeline_kwargs = config.eval.get_pipeline_kwargs()
29
30 task = config.pipeline.task
31 control_root = config.eval.control_root
32 if task in ["canny-to-image", "depth-to-image", "inpainting"]:
33 controls = get_control(
34 task,
35 batch["image"],
36 names=batch["filename"],
37 data_root=os.path.join(
38 control_root, collect_config.dataset_name, f"{dataset.config_name}-{config.eval.num_samples}"
39 ),
40 )
41 if task == "inpainting":
42 pipeline_kwargs["image"] = controls[0]
43 pipeline_kwargs["mask_image"] = controls[1]
44 else:
45 pipeline_kwargs["control_image"] = controls
46
47 # Handle meta tensors by moving individual components
48 try:
49 pipeline = pipeline.to("cuda")
50 except NotImplementedError:
51 # Move individual pipeline components that have to_empty method
52 if hasattr(pipeline, 'transformer') and pipeline.transformer is not None:
53 try:
54 pipeline.transformer = pipeline.transformer.to("cuda")
55 except NotImplementedError:
56 pipeline.transformer = pipeline.transformer.to_empty(device="cuda")
57
58 if hasattr(pipeline, 'text_encoder') and pipeline.text_encoder is not None:
59 try:
60 pipeline.text_encoder = pipeline.text_encoder.to("cuda")
61 except NotImplementedError:
62 pipeline.text_encoder = pipeline.text_encoder.to_empty(device="cuda")
63
64 if hasattr(pipeline, 'text_encoder_2') and pipeline.text_encoder_2 is not None:
65 try:
66 pipeline.text_encoder_2 = pipeline.text_encoder_2.to("cuda")
67 except NotImplementedError:
68 pipeline.text_encoder_2 = pipeline.text_encoder_2.to_empty(device="cuda")
69
70 if hasattr(pipeline, 'vae') and pipeline.vae is not None:
71 try:
72 pipeline.vae = pipeline.vae.to("cuda")
73 except NotImplementedError:
74 pipeline.vae = pipeline.vae.to_empty(device="cuda")
75
76 result_images = pipeline(prompt=prompts, generator=generators, **pipeline_kwargs).images
77 num_guidances = (len(caches) // batch_size) // config.eval.num_steps
78 num_steps = len(caches) // (batch_size * num_guidances)
79 assert (
80 len(caches) == batch_size * num_steps * num_guidances
81 ), f"Unexpected number of caches: {len(caches)} != {batch_size} * {config.eval.num_steps} * {num_guidances}"
82 for j, (filename, image) in enumerate(zip(filenames, result_images, strict=True)):
83 image.save(os.path.join(samples_dirpath, f"{filename}.png"))
84 for s in range(num_steps):
85 for g in range(num_guidances):
86 c = caches[s * batch_size * num_guidances + g * batch_size + j]
87 c["filename"] = filename
88 c["step"] = s
89 c["guidance"] = g
90 c = tree_map(lambda x: process(x), c)
91 torch.save(c, os.path.join(caches_dirpath, f"{filename}-{s:05d}-{g}.pt"))
92 caches.clear()1def quantize_scale(
2 s: torch.Tensor,
3 /,
4 *,
5 quant_dtypes: tp.Sequence[QuantDataType],
6 quant_spans: tp.Sequence[float],
7 view_shapes: tp.Sequence[torch.Size],
8) -> QuantScale:
9 """Quantize the scale tensor.
10
11 Args:
12 s (`torch.Tensor`):
13 The scale tensor.
14 quant_dtypes (`Sequence[QuantDataType]`):
15 The quantization dtypes of the scale tensor.
16 quant_spans (`Sequence[float]`):
17 The quantization spans of the scale tensor.
18 view_shapes (`Sequence[torch.Size]`):
19 The view shapes of the scale tensor.
20
21 Returns:
22 `QuantScale`:
23 The quantized scale tensor.
24 """
25 # Add validation at the start
26 if s.numel() == 0:
27 raise ValueError("Input tensor is empty")
28 if s.isnan().any() or s.isinf().any():
29 raise ValueError("Input tensor contains NaN or Inf values")
30 if (s == 0).all():
31 raise ValueError("Input tensor contains all zeros")
32
33 # Add meta tensor check before any operations
34 if s.is_meta:
35 raise RuntimeError("Cannot quantize scale with meta tensor. Ensure model is loaded on actual device.")
36
37 # Existing validation
38 if s.isnan().any() or s.isinf().any():
39 raise ValueError("Input tensor contains NaN or Inf values")
40
41 scale = QuantScale()
42 s = s.abs()
43 for view_shape, quant_dtype, quant_span in zip(view_shapes[:-1], quant_dtypes[:-1], quant_spans[:-1], strict=True):
44 s = s.view(view_shape) # (#g0, rs0, #g1, rs1, #g2, rs2, ...)
45 ss = s.amax(dim=list(range(1, len(view_shape), 2)), keepdim=True) # i.e., s_dynamic_span
46 ss = simple_quantize(
47 ss / quant_span, has_zero_point=False, quant_dtype=quant_dtype
48 ) # i.e., s_scale = s_dynamic_span / s_quant_span
49 s = s / ss
50 scale.append(ss)
51 view_shape = view_shapes[-1]
52 s = s.view(view_shape)
53 if any(v != 1 for v in view_shape[1::2]):
54 ss = s.amax(dim=list(range(1, len(view_shape), 2)), keepdim=True)
55 ss = simple_quantize(ss / quant_spans[-1], has_zero_point=False, quant_dtype=quant_dtypes[-1])
56 else:
57 assert quant_spans[-1] == 1, "The last quant span must be 1."
58 ss = simple_quantize(s, has_zero_point=False, quant_dtype=quant_dtypes[-1])
59 scale.append(ss)
60 scale.remove_zero()
61 return scale
62
63 def quantize(
64 self,
65 *,
66 # scale-based quantization related arguments
67 scale: torch.Tensor | None = None,
68 zero: torch.Tensor | None = None,
69 # range-based quantization related arguments
70 tensor: torch.Tensor | None = None,
71 dynamic_range: DynamicRange | None = None,
72 ) -> tuple[QuantScale, torch.Tensor]:
73 """Get the quantization scale and zero point of the tensor to be quantized.
74
75 Args:
76 scale (`torch.Tensor` or `None`, *optional*, defaults to `None`):
77 The scale tensor.
78 zero (`torch.Tensor` or `None`, *optional*, defaults to `None`):
79 The zero point tensor.
80 tensor (`torch.Tensor` or `None`, *optional*, defaults to `None`):
81 Ten tensor to be quantized. This is only used for range-based quantization.
82 dynamic_range (`DynamicRange` or `None`, *optional*, defaults to `None`):
83 The dynamic range of the tensor to be quantized.
84
85 Returns:
86 `tuple[QuantScale, torch.Tensor]`:
87 The scale and the zero point.
88 """
89 # region step 1: get the dynamic span for range-based scale or the scale tensor
90 if scale is None:
91 range_based = True
92 assert isinstance(tensor, torch.Tensor), "View tensor must be a tensor."
93 dynamic_range = dynamic_range or DynamicRange()
94 dynamic_range = dynamic_range.measure(
95 tensor.view(self.tensor_view_shape),
96 zero_domain=self.tensor_zero_domain,
97 is_float_point=self.tensor_quant_dtype.is_float_point,
98 )
99 dynamic_range = dynamic_range.intersect(self.tensor_range_bound)
100 dynamic_span = (dynamic_range.max - dynamic_range.min) if self.has_zero_point else dynamic_range.max
101 else:
102 range_based = False
103 scale = scale.view(self.scale_view_shapes[-1])
104 assert isinstance(scale, torch.Tensor), "Scale must be a tensor."
105 # endregion
106 # region step 2: get the scale
107 if self.linear_scale_quant_dtypes:
108 if range_based:
109 linear_scale = dynamic_span / self.linear_tensor_quant_span
110 elif self.exponent_scale_quant_dtypes:
111 linear_scale = scale.mul(self.exponent_tensor_quant_span).div(self.linear_tensor_quant_span)
112 else:
113 linear_scale = scale
114 lin_s = quantize_scale(
115 linear_scale,
116 quant_dtypes=self.linear_scale_quant_dtypes,
117 quant_spans=self.linear_scale_quant_spans,
118 view_shapes=self.linear_scale_view_shapes,
119 )
120 assert lin_s.data is not None, "Linear scale tensor is None."
121 if not lin_s.data.is_meta:
122 assert not lin_s.data.isnan().any(), "Linear scale tensor contains NaN."
123 assert not lin_s.data.isinf().any(), "Linear scale tensor contains Inf."
124 else:
125 lin_s = QuantScale()
126 if self.exponent_scale_quant_dtypes:
127 if range_based:
128 exp_scale = dynamic_span / self.exponent_tensor_quant_span
129 else:
130 exp_scale = scale
131 if lin_s.data is not None:
132 lin_s.data = lin_s.data.expand(self.linear_scale_view_shapes[-1]).reshape(self.scale_view_shapes[-1])
133 exp_scale = exp_scale / lin_s.data
134 exp_s = quantize_scale(
135 exp_scale,
136 quant_dtypes=self.exponent_scale_quant_dtypes,
137 quant_spans=self.exponent_scale_quant_spans,
138 view_shapes=self.exponent_scale_view_shapes,
139 )
140 assert exp_s.data is not None, "Exponential scale tensor is None."
141 assert not exp_s.data.isnan().any(), "Exponential scale tensor contains NaN."
142 assert not exp_s.data.isinf().any(), "Exponential scale tensor contains Inf."
143 s = exp_s if lin_s.data is None else lin_s.extend(exp_s)
144 else:
145 s = lin_s
146
147 # Before the final assertions, add debugging and validation
148 if s.data is None:
149 # Log debugging information
150 print(f"Linear scale dtypes: {self.linear_scale_quant_dtypes}")
151 print(f"Exponent scale dtypes: {self.exponent_scale_quant_dtypes}")
152 if hasattr(lin_s, 'data') and lin_s.data is not None:
153 print(f"Linear scale data shape: {lin_s.data.shape}")
154 raise RuntimeError("Scale computation failed - resulting scale is None")
155 assert s.data is not None, "Scale tensor is None."
156 assert not s.data.isnan().any(), "Scale tensor contains NaN."
157 assert not s.data.isinf().any(), "Scale tensor contains Inf."
158 # endregion
159 # region step 3: get the zero point
160 if self.has_zero_point:
161 if range_based:
162 if self.tensor_zero_domain == ZeroPointDomain.PreScale:
163 zero = self.tensor_quant_range.min - dynamic_range.min / s.data
164 else:
165 zero = self.tensor_quant_range.min * s.data - dynamic_range.min
166 assert isinstance(zero, torch.Tensor), "Zero point must be a tensor."
167 z = simple_quantize(zero, has_zero_point=True, quant_dtype=self.zero_quant_dtype)
168 else:
169 z = torch.tensor(0, dtype=s.data.dtype, device=s.data.device)
170 assert not z.isnan().any(), "Zero point tensor contains NaN."
171 assert not z.isinf().any(), "Zero point tensor contains Inf."
172 # endregion
173 return s, z1def ptq( # noqa: C901
2 model: DiffusionModelStruct,
3 config: DiffusionQuantConfig,
4 cache: DiffusionPtqCacheConfig | None = None,
5 load_dirpath: str = "",
6 save_dirpath: str = "",
7 copy_on_save: bool = False,
8 save_model: bool = False,
9) -> DiffusionModelStruct:
10 """Post-training quantization of a diffusion model.
11
12 Args:
13 model (`DiffusionModelStruct`):
14 The diffusion model.
15 config (`DiffusionQuantConfig`):
16 The diffusion model post-training quantization configuration.
17 cache (`DiffusionPtqCacheConfig`, *optional*, defaults to `None`):
18 The diffusion model quantization cache path configuration.
19 load_dirpath (`str`, *optional*, defaults to `""`):
20 The directory path to load the quantization checkpoint.
21 save_dirpath (`str`, *optional*, defaults to `""`):
22 The directory path to save the quantization checkpoint.
23 copy_on_save (`bool`, *optional*, defaults to `False`):
24 Whether to copy the cache to the save directory.
25 save_model (`bool`, *optional*, defaults to `False`):
26 Whether to save the quantized model checkpoint.
27
28 Returns:
29 `DiffusionModelStruct`:
30 The quantized diffusion model.
31 """
32 logger = tools.logging.getLogger(__name__)
33 if not isinstance(model, DiffusionModelStruct):
34 model = DiffusionModelStruct.construct(model)
35 assert isinstance(model, DiffusionModelStruct)
36
37 quant_wgts = config.enabled_wgts
38 quant_ipts = config.enabled_ipts
39 quant_opts = config.enabled_opts
40 quant_acts = quant_ipts or quant_opts
41 quant = quant_wgts or quant_acts
42
43 load_model_path, load_path, save_path = "", None, None
44 if load_dirpath:
45 load_path = DiffusionQuantCacheConfig(
46 smooth=os.path.join(load_dirpath, "smooth.pt"),
47 branch=os.path.join(load_dirpath, "branch.pt"),
48 wgts=os.path.join(load_dirpath, "wgts.pt"),
49 acts=os.path.join(load_dirpath, "acts.pt"),
50 )
51 load_model_path = os.path.join(load_dirpath, "model.pt")
52 if os.path.exists(load_model_path):
53 if config.enabled_wgts and config.wgts.enabled_low_rank:
54 if os.path.exists(load_path.branch):
55 load_model = True
56 else:
57 logger.warning(f"Model low-rank branch checkpoint {load_path.branch} does not exist")
58 load_model = False
59 else:
60 load_model = True
61 if load_model:
62 logger.info(f"* Loading model from {load_model_path}")
63 save_dirpath = "" # do not save the model if loading
64 else:
65 logger.warning(f"Model checkpoint {load_model_path} does not exist")
66 load_model = False
67 else:
68 load_model = False
69 if save_dirpath:
70 os.makedirs(save_dirpath, exist_ok=True)
71 save_path = DiffusionQuantCacheConfig(
72 smooth=os.path.join(save_dirpath, "smooth.pt"),
73 branch=os.path.join(save_dirpath, "branch.pt"),
74 wgts=os.path.join(save_dirpath, "wgts.pt"),
75 acts=os.path.join(save_dirpath, "acts.pt"),
76 )
77 else:
78 save_model = False
79
80 if quant and config.enabled_rotation:
81 logger.info("* Rotating model for quantization")
82 tools.logging.Formatter.indent_inc()
83 rotate_diffusion(model, config=config)
84 tools.logging.Formatter.indent_dec()
85 gc.collect()
86 torch.cuda.empty_cache()
87
88 # region smooth quantization
89 if quant and config.enabled_smooth:
90 logger.info("* Smoothing model for quantization")
91 tools.logging.Formatter.indent_inc()
92 load_from = ""
93 if load_path and os.path.exists(load_path.smooth):
94 load_from = load_path.smooth
95 elif cache and cache.path.smooth and os.path.exists(cache.path.smooth):
96 load_from = cache.path.smooth
97 if load_from:
98 logger.info(f"- Loading smooth scales from {load_from}")
99 smooth_cache = torch.load(load_from)
100 smooth_diffusion(model, config, smooth_cache=smooth_cache)
101 else:
102 logger.info("- Generating smooth scales")
103 smooth_cache = smooth_diffusion(model, config)
104 if cache and cache.path.smooth:
105 logger.info(f"- Saving smooth scales to {cache.path.smooth}")
106 os.makedirs(cache.dirpath.smooth, exist_ok=True)
107 torch.save(smooth_cache, cache.path.smooth)
108 load_from = cache.path.smooth
109 if save_path:
110 if not copy_on_save and load_from:
111 logger.info(f"- Linking smooth scales to {save_path.smooth}")
112 os.symlink(os.path.relpath(load_from, save_dirpath), save_path.smooth)
113 else:
114 logger.info(f"- Saving smooth scales to {save_path.smooth}")
115 torch.save(smooth_cache, save_path.smooth)
116 del smooth_cache
117 tools.logging.Formatter.indent_dec()
118 gc.collect()
119 torch.cuda.empty_cache()
120 # endregion
121 # region collect original state dict
122 if config.needs_acts_quantizer_cache:
123 if load_path and os.path.exists(load_path.acts):
124 orig_state_dict = None
125 elif cache and cache.path.acts and os.path.exists(cache.path.acts):
126 orig_state_dict = None
127 else:
128 orig_state_dict: dict[str, torch.Tensor] = {
129 name: param.detach().clone() for name, param in model.module.named_parameters() if param.ndim > 1
130 }
131 else:
132 orig_state_dict = None
133 # endregion
134 if load_model:
135 logger.info(f"* Loading model checkpoint from {load_model_path}")
136 load_diffusion_weights_state_dict(
137 model,
138 config,
139 state_dict=torch.load(load_model_path),
140 branch_state_dict=torch.load(load_path.branch) if os.path.exists(load_path.branch) else None,
141 )
142 gc.collect()
143 torch.cuda.empty_cache()
144 elif quant_wgts:
145 logger.info("* Ensuring model is on actual device before quantization")
146
147 # Check if model has meta tensors
148 has_meta_tensors = any(param.is_meta for param in model.module.parameters())
149
150 if has_meta_tensors:
151 logger.info("* Model contains meta tensors, materializing to actual device")
152
153 # Option 1: Use to_empty() and reload weights (recommended)
154 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
155
156 # Store original state dict if available
157 try:
158 original_state_dict = model.module.state_dict()
159 model.module = model.module.to_empty(device=device)
160 model.module.load_state_dict(original_state_dict)
161 logger.info("* Successfully materialized model with original weights")
162 except Exception as e:
163 logger.warning(f"* Failed to preserve weights during materialization: {e}")
164 # Fallback: just move to empty device (weights will be zero)
165 model.module = model.module.to_empty(device=device)
166 logger.warning("* Model moved to device but weights may be uninitialized")
167 else:
168 # Model already has real tensors, just ensure it's on the right device
169 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
170 model.module = model.module.to(device)
171
172 # Verify no meta tensors remain
173 remaining_meta = [name for name, param in model.module.named_parameters() if param.is_meta]
174 if remaining_meta:
175 raise RuntimeError(f"Parameters still on meta device: {remaining_meta}")
176
177 logger.info("* Model successfully prepared for quantization")
178
179 logger.info("* Quantizing weights")
180 tools.logging.Formatter.indent_inc()
181 quantizer_state_dict, quantizer_load_from = None, ""
182 if load_path and os.path.exists(load_path.wgts):
183 quantizer_load_from = load_path.wgts
184 elif cache and cache.path.wgts and os.path.exists(cache.path.wgts):
185 quantizer_load_from = cache.path.wgts
186 if quantizer_load_from:
187 logger.info(f"- Loading weight settings from {quantizer_load_from}")
188 quantizer_state_dict = torch.load(quantizer_load_from)
189 branch_state_dict, branch_load_from = None, ""
190 if load_path and os.path.exists(load_path.branch):
191 branch_load_from = load_path.branch
192 elif cache and cache.path.branch and os.path.exists(cache.path.branch):
193 branch_load_from = cache.path.branch
194 if branch_load_from:
195 logger.info(f"- Loading branch settings from {branch_load_from}")
196 branch_state_dict = torch.load(branch_load_from)
197 if not quantizer_load_from:
198 logger.info("- Generating weight settings")
199 if not branch_load_from:
200 logger.info("- Generating branch settings")
201 quantizer_state_dict, branch_state_dict, scale_state_dict = quantize_diffusion_weights(
202 model,
203 config,
204 quantizer_state_dict=quantizer_state_dict,
205 branch_state_dict=branch_state_dict,
206 return_with_scale_state_dict=bool(save_dirpath),
207 )
208 if not quantizer_load_from and cache and cache.dirpath.wgts:
209 logger.info(f"- Saving weight settings to {cache.path.wgts}")
210 os.makedirs(cache.dirpath.wgts, exist_ok=True)
211 torch.save(quantizer_state_dict, cache.path.wgts)
212 quantizer_load_from = cache.path.wgts
213 if not branch_load_from and cache and cache.dirpath.branch:
214 logger.info(f"- Saving branch settings to {cache.path.branch}")
215 os.makedirs(cache.dirpath.branch, exist_ok=True)
216 torch.save(branch_state_dict, cache.path.branch)
217 branch_load_from = cache.path.branch
218 if save_path:
219 if not copy_on_save and quantizer_load_from:
220 logger.info(f"- Linking weight settings to {save_path.wgts}")
221 os.symlink(os.path.relpath(quantizer_load_from, save_dirpath), save_path.wgts)
222 else:
223 logger.info(f"- Saving weight settings to {save_path.wgts}")
224 torch.save(quantizer_state_dict, save_path.wgts)
225 if not copy_on_save and branch_load_from:
226 logger.info(f"- Linking branch settings to {save_path.branch}")
227 os.symlink(os.path.relpath(branch_load_from, save_dirpath), save_path.branch)
228 else:
229 logger.info(f"- Saving branch settings to {save_path.branch}")
230 torch.save(branch_state_dict, save_path.branch)
231 if save_model:
232 logger.info(f"- Saving model to {save_dirpath}")
233 torch.save(scale_state_dict, os.path.join(save_dirpath, "scale.pt"))
234 torch.save(model.module.state_dict(), os.path.join(save_dirpath, "model.pt"))
235 del quantizer_state_dict, branch_state_dict, scale_state_dict
236 tools.logging.Formatter.indent_dec()
237 gc.collect()
238 torch.cuda.empty_cache()
239 if quant_acts:
240 logger.info(" * Quantizing activations")
241 tools.logging.Formatter.indent_inc()
242 if config.needs_acts_quantizer_cache:
243 load_from = ""
244 if load_path and os.path.exists(load_path.acts):
245 load_from = load_path.acts
246 elif cache and cache.path.acts and os.path.exists(cache.path.acts):
247 load_from = cache.path.acts
248 if load_from:
249 logger.info(f"- Loading activation settings from {load_from}")
250 quantizer_state_dict = torch.load(load_from)
251 quantize_diffusion_activations(
252 model, config, quantizer_state_dict=quantizer_state_dict, orig_state_dict=orig_state_dict
253 )
254 else:
255 logger.info("- Generating activation settings")
256 quantizer_state_dict = quantize_diffusion_activations(model, config, orig_state_dict=orig_state_dict)
257 if cache and cache.dirpath.acts and quantizer_state_dict is not None:
258 logger.info(f"- Saving activation settings to {cache.path.acts}")
259 os.makedirs(cache.dirpath.acts, exist_ok=True)
260 torch.save(quantizer_state_dict, cache.path.acts)
261 load_from = cache.path.acts
262 if save_dirpath:
263 if not copy_on_save and load_from:
264 logger.info(f"- Linking activation quantizer settings to {save_path.acts}")
265 os.symlink(os.path.relpath(load_from, save_dirpath), save_path.acts)
266 else:
267 logger.info(f"- Saving activation quantizer settings to {save_path.acts}")
268 torch.save(quantizer_state_dict, save_path.acts)
269 del quantizer_state_dict
270 else:
271 logger.info("- No need to generate/load activation quantizer settings")
272 quantize_diffusion_activations(model, config, orig_state_dict=orig_state_dict)
273 tools.logging.Formatter.indent_dec()
274 del orig_state_dict
275 gc.collect()
276 torch.cuda.empty_cache()
277 return model