33 lines
2.7 KiB
Python
33 lines
2.7 KiB
Python
"""Probe loaded Comfy H3 patch projections on their captured FP32 rows."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
|
source = path.read_text(encoding="utf-8")
|
|
if "import os\n" not in source:
|
|
source = source.replace("import math\n", "import math\nimport os\n", 1)
|
|
old = (
|
|
" video_embed = self.video_patch_proj(all_video_rows).to(dtype)\n"
|
|
" audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)\n"
|
|
)
|
|
new = (
|
|
" video_embed_fp32 = self.video_patch_proj(all_video_rows)\n"
|
|
" audio_embed_fp32 = self.audio_patch_proj(all_audio_rows)\n"
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
|
|
" if capture_dir and not os.path.exists(os.path.join(capture_dir, \"h3_loaded_patch_projection_probe.pt\")):\n"
|
|
" def probe(layer, rows, output):\n"
|
|
" weight, bias, state = comfy.ops.cast_bias_weight(layer, rows, offloadable=True)\n"
|
|
" try:\n"
|
|
" cast_output = torch.nn.functional.linear(rows, weight, bias)\n"
|
|
" return {\"rows\": rows.detach().cpu(), \"output\": output.detach().cpu(), \"forward_cast\": layer.forward_comfy_cast_weights(rows).detach().cpu(), \"effective_weight\": weight.detach().cpu(), \"effective_bias\": None if bias is None else bias.detach().cpu(), \"effective_weight_dtype\": str(weight.dtype), \"effective_bias_dtype\": None if bias is None else str(bias.dtype), \"cast_output\": cast_output.detach().cpu(), \"rows_dtype\": str(rows.dtype), \"rows_stride\": tuple(rows.stride()), \"weight_dtype\": str(layer.weight.dtype), \"weight_stride\": tuple(layer.weight.stride()), \"bias_dtype\": None if layer.bias is None else str(layer.bias.dtype), \"bias_stride\": None if layer.bias is None else tuple(layer.bias.stride()), \"autocast\": torch.is_autocast_enabled(), \"force_cast\": getattr(layer, \"comfy_force_cast_weights\", False), \"weight_functions\": len(layer.weight_function), \"bias_functions\": len(layer.bias_function)}\n"
|
|
" finally:\n"
|
|
" comfy.ops.uncast_bias_weight(layer, weight, bias, state)\n"
|
|
" torch.save({\"video\": probe(self.video_patch_proj, all_video_rows, video_embed_fp32), \"audio\": probe(self.audio_patch_proj, all_audio_rows, audio_embed_fp32)}, os.path.join(capture_dir, \"h3_loaded_patch_projection_probe.pt\"))\n"
|
|
" video_embed = video_embed_fp32.to(dtype)\n"
|
|
" audio_embed = audio_embed_fp32.to(dtype)\n"
|
|
)
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate H3 patch-projection calls.")
|
|
path.write_text(source.replace(old, new), encoding="utf-8")
|
|
print("Applied loaded H3 patch-projection probe.")
|