"""Capture loaded Comfy Qwen QuantizedTensor fields and dispatch variants."""

from pathlib import Path


path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
source = path.read_text(encoding="utf-8")
helper = (
    "def _h3_owned_cpu(tensor):\n"
    "    return tensor.detach().to(device=\"cpu\", copy=True).contiguous().clone()\n\n"
)
if "def _h3_owned_cpu(tensor):\n" not in source:
    source = helper + source
old = "        xq = self.q_proj(hidden_states)\n        xk = self.k_proj(hidden_states)\n        xv = self.v_proj(hidden_states)\n"
new = old + (
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
    "        if capture_dir:\n"
    "            modules = {\"q\": self.q_proj, \"k\": self.k_proj, \"v\": self.v_proj}\n"
    "            def probe(module):\n"
    "                stored = module.weight\n"
    "                qdata, scale, block_scale = stored.layout_cls.get_plain_tensors(stored)\n"
    "                weight, bias, offload_stream = comfy.ops.cast_bias_weight(module, hidden_states, offloadable=True)\n"
    "                try:\n"
    "                    output = torch.nn.functional.linear(hidden_states, weight, bias)\n"
    "                    return {\"qdata\": _h3_owned_cpu(qdata), \"scale\": _h3_owned_cpu(scale), \"block_scale\": _h3_owned_cpu(block_scale), \"layout_class\": stored._layout_cls, \"orig_dtype\": stored._params.orig_dtype, \"orig_shape\": tuple(stored._params.orig_shape), \"effective_weight\": _h3_owned_cpu(weight), \"effective_dtype\": str(weight.dtype), \"effective_shape\": tuple(weight.shape), \"output\": _h3_owned_cpu(output)}\n"
    "                finally:\n"
    "                    comfy.ops.uncast_bias_weight(module, weight, bias, offload_stream)\n"
    "            torch.save({\"input\": hidden_states.detach().cpu(), \"modules\": {name: probe(module) for name, module in modules.items()}}, os.path.join(capture_dir, \"qwen0_quantized_qkv.pt\"))\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate Qwen QKV projections.")
source = source.replace(old, new)

old = "        return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))\n"
new = (
    "        gate = self.gate_proj(x)\n"
    "        up = self.up_proj(x)\n"
    "        activated = self.activation(gate) * up\n"
    "        output = self.down_proj(activated)\n"
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
    "        if capture_dir:\n"
    "            modules = {\"gate\": self.gate_proj, \"up\": self.up_proj, \"down\": self.down_proj}\n"
    "            def probe(module, value):\n"
    "                stored = module.weight\n"
    "                qdata, scale, block_scale = stored.layout_cls.get_plain_tensors(stored)\n"
    "                weight, bias, offload_stream = comfy.ops.cast_bias_weight(module, value, offloadable=True)\n"
    "                try:\n"
    "                    output = torch.nn.functional.linear(value, weight, bias)\n"
    "                    return {\"qdata\": _h3_owned_cpu(qdata), \"scale\": _h3_owned_cpu(scale), \"block_scale\": _h3_owned_cpu(block_scale), \"layout_class\": stored._layout_cls, \"orig_dtype\": stored._params.orig_dtype, \"orig_shape\": tuple(stored._params.orig_shape), \"effective_weight\": _h3_owned_cpu(weight), \"effective_dtype\": str(weight.dtype), \"effective_shape\": tuple(weight.shape), \"output\": _h3_owned_cpu(output)}\n"
    "                finally:\n"
    "                    comfy.ops.uncast_bias_weight(module, weight, bias, offload_stream)\n"
    "            torch.save({\"input\": x.detach().cpu(), \"activated\": activated.detach().cpu(), \"modules\": {\"gate\": probe(self.gate_proj, x), \"up\": probe(self.up_proj, x), \"down\": probe(self.down_proj, activated)}}, os.path.join(capture_dir, \"qwen0_quantized_mlp.pt\"))\n"
    "        return output\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate Qwen MLP projections.")
path.write_text(source.replace(old, new), encoding="utf-8")
print("Applied Qwen QuantizedTensor dispatch probe.")