"""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.")