75 lines
5.5 KiB
Python
75 lines
5.5 KiB
Python
|
"""Capture actual loaded-Comfy Qwen merged and DeepStack vision tensors."""
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
path = Path("/opt/ComfyUI/comfy/text_encoders/qwen3vl.py")
|
||
|
|
source = path.read_text(encoding="utf-8")
|
||
|
|
if "import os\n" not in source:
|
||
|
|
source = source.replace("import os\n", "import os\n", 1) if "import os\n" in source else "import os\n" + source
|
||
|
|
if "import traceback\n" not in source:
|
||
|
|
source = "import traceback\n" + source
|
||
|
|
old = (
|
||
|
|
" merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid)\n"
|
||
|
|
" return merged, {\"grid\": grid, \"deepstack\": deepstack}\n"
|
||
|
|
)
|
||
|
|
new = (
|
||
|
|
" merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid)\n"
|
||
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
|
||
|
|
" if capture_dir:\n"
|
||
|
|
" os.makedirs(capture_dir, exist_ok=True)\n"
|
||
|
|
" capture_index = getattr(self, \"_h3_vision_capture_index\", 0)\n"
|
||
|
|
" torch.save({\"merged\": merged.detach().cpu(), \"deepstack\": [value.detach().cpu() for value in deepstack], \"pixel_values\": image.detach().cpu(), \"grid\": grid.detach().cpu(), \"stack\": traceback.format_stack()}, os.path.join(capture_dir, f\"qwen_vision_{capture_index}.pt\"))\n"
|
||
|
|
" open(os.path.join(capture_dir, \"qwen_vision_stack.txt\"), \"w\", encoding=\"utf-8\").writelines(traceback.format_stack())\n"
|
||
|
|
" self._h3_vision_capture_index = capture_index + 1\n"
|
||
|
|
" return merged, {\"grid\": grid, \"deepstack\": deepstack}\n"
|
||
|
|
)
|
||
|
|
if source.count(old) != 1:
|
||
|
|
raise RuntimeError("Unable to locate Qwen3-VL vision preprocess block.")
|
||
|
|
path.write_text(source.replace(old, new), encoding="utf-8")
|
||
|
|
|
||
|
|
vision_path = Path("/opt/ComfyUI/comfy/text_encoders/qwen35.py")
|
||
|
|
vision_source = vision_path.read_text(encoding="utf-8")
|
||
|
|
vision_source = vision_source.replace(
|
||
|
|
" x = self.patch_embed(x)\n pos_embeds = self.fast_pos_embed_interpolate(grid_thw).to(x.device)\n x = x + pos_embeds\n",
|
||
|
|
" x = self.patch_embed(x)\n capture_trace = {\"patch_embed\": x.detach().cpu()}\n pos_embeds = self.fast_pos_embed_interpolate(grid_thw).to(x.device)\n capture_trace[\"position_embed\"] = pos_embeds.detach().cpu()\n x = x + pos_embeds\n capture_trace[\"vision_input\"] = x.detach().cpu()\n",
|
||
|
|
1,
|
||
|
|
)
|
||
|
|
vision_source = vision_source.replace(
|
||
|
|
" x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, optimized_attention=optimized_attention)\n",
|
||
|
|
" x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, optimized_attention=optimized_attention)\n capture_trace[f\"block_{layer_num:02d}\"] = x.detach().cpu()\n",
|
||
|
|
1,
|
||
|
|
)
|
||
|
|
vision_source = vision_source.replace(
|
||
|
|
" merged = self.merger(x)\n if self.deepstack_merger_list is not None:\n",
|
||
|
|
" merged = self.merger(x)\n capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n if capture_dir:\n trace_index = getattr(self, \"_h3_vision_trace_index\", 0)\n torch.save(capture_trace, os.path.join(capture_dir, f\"qwen_vision_trace_{trace_index}.pt\"))\n self._h3_vision_trace_index = trace_index + 1\n if self.deepstack_merger_list is not None:\n",
|
||
|
|
1,
|
||
|
|
)
|
||
|
|
vision_path.write_text(vision_source, encoding="utf-8")
|
||
|
|
|
||
|
|
clip_path = Path("/opt/ComfyUI/comfy/sd1_clip.py")
|
||
|
|
clip_source = clip_path.read_text(encoding="utf-8")
|
||
|
|
if "import os\n" not in clip_source:
|
||
|
|
clip_source = "import os\n" + clip_source
|
||
|
|
clip_source = clip_source.replace(
|
||
|
|
" tokens_embed = torch.tensor([tokens_temp], device=device, dtype=torch.long)\n tokens_embed = self.transformer.get_input_embeddings()(tokens_embed, out_dtype=torch.float32)\n",
|
||
|
|
" tokens_embed = torch.tensor([tokens_temp], device=device, dtype=torch.long)\n capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n if capture_dir:\n torch.save(tokens_embed.detach().cpu(), os.path.join(capture_dir, \"qwen_compact_token_ids.pt\"))\n tokens_embed = self.transformer.get_input_embeddings()(tokens_embed, out_dtype=torch.float32)\n",
|
||
|
|
1,
|
||
|
|
)
|
||
|
|
clip_path.write_text(clip_source, encoding="utf-8")
|
||
|
|
|
||
|
|
nodes_path = Path("/opt/ComfyUI/comfy_extras/nodes_minimax_h3.py")
|
||
|
|
nodes_source = nodes_path.read_text(encoding="utf-8")
|
||
|
|
if "import os\n" not in nodes_source:
|
||
|
|
nodes_source = "import os\n" + nodes_source
|
||
|
|
vae_encode_line = "kf[\"latent\"] = vae.encode(kf.pop(\"image\"))"
|
||
|
|
if nodes_source.count(vae_encode_line) != 1:
|
||
|
|
raise RuntimeError("Unable to locate MiniMax H3 keyframe VAE encode call.")
|
||
|
|
nodes_source = nodes_source.replace(
|
||
|
|
vae_encode_line,
|
||
|
|
"keyframe_image = kf.pop(\"image\")\n kf[\"latent\"] = vae.encode(keyframe_image)\n capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n if capture_dir:\n capture_index = getattr(vae, \"_h3_vae_capture_index\", 0)\n vae_model = vae.first_stage_model\n vae_meta = {\"parameter_dtype\": str(next(vae_model.parameters()).dtype), \"tiling\": vae_model.tiling, \"tile_size\": vae_model.tile_size, \"tile_overlap_min\": vae_model.tile_overlap_min}\n torch.save({\"image\": keyframe_image.detach().cpu(), \"latent\": kf[\"latent\"].detach().cpu(), \"meta\": vae_meta}, os.path.join(capture_dir, f\"vae_keyframe_{capture_index}.pt\"))\n vae._h3_vae_capture_index = capture_index + 1",
|
||
|
|
1,
|
||
|
|
)
|
||
|
|
nodes_path.write_text(nodes_source, encoding="utf-8")
|
||
|
|
print("Applied Qwen3-VL vision capture patch.")
|