after add cache, this version support flashinfer

configs/hunyuan_image3/hunyuan_image3_i2i.json

@@ -0,0 +1,28 @@
+{
+    "infer_steps": 50,
+    "sample_guide_scale": 5.0,
+    "flow_shift": 1.0,
+    "image_size": "auto",
+    "feature_caching": "NoCaching",
+    "enable_kv_cache": true,
+    "enable_text_kv_cache": true,
+    "use_taylor_cache": false,
+    "taylor_cache_interval": 5,
+    "taylor_cache_order": 2,
+    "taylor_cache_enable_first_enhance": false,
+    "taylor_cache_first_enhance_steps": 3,
+    "taylor_cache_enable_tailing_enhance": false,
+    "taylor_cache_tailing_enhance_steps": 1,
+    "taylor_cache_low_freqs_order": 2,
+    "taylor_cache_high_freqs_order": 2,
+    "pipeline_parallel": true,
+    "enable_cfg": true,
+    "infer_align_image_size": true,
+    "bot_task": "think_recaption",
+    "use_system_prompt": "en_unified",
+    "max_new_tokens": 2048,
+    "text_do_sample": true,
+    "text_top_k": 1024,
+    "text_top_p": 0.95,
+    "text_temperature": 0.6
+}

configs/hunyuan_image3/hunyuan_image3_t2i.json

@@ -0,0 +1,28 @@
+{
+    "infer_steps": 50,
+    "sample_guide_scale": 5.0,
+    "flow_shift": 1.0,
+    "target_height": 1024,
+    "target_width": 1024,
+    "feature_caching": "NoCaching",
+    "enable_kv_cache": true,
+    "enable_text_kv_cache": true,
+    "use_taylor_cache": false,
+    "taylor_cache_interval": 5,
+    "taylor_cache_order": 2,
+    "taylor_cache_enable_first_enhance": false,
+    "taylor_cache_first_enhance_steps": 3,
+    "taylor_cache_enable_tailing_enhance": false,
+    "taylor_cache_tailing_enhance_steps": 1,
+    "taylor_cache_low_freqs_order": 2,
+    "taylor_cache_high_freqs_order": 2,
+    "pipeline_parallel": true,
+    "enable_cfg": true,
+    "bot_task": "think_recaption",
+    "use_system_prompt": "en_unified",
+    "max_new_tokens": 2048,
+    "text_do_sample": true,
+    "text_top_k": 1024,
+    "text_top_p": 0.95,
+    "text_temperature": 0.6
+}

docs/superpowers/plans/2026-07-02-hunyuan-image3-think-recaption.md

@@ -0,0 +1,45 @@
+# Hunyuan Image3 Think Recaption Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Add a native LightX2V text-generation pre-stage for Hunyuan Image3 `recaption` and `think_recaption`, then feed generated COT text into the existing native image generation path.
+
+**Architecture:** Keep the DiT path native. Add small runner-level helpers for text input preparation, autoregressive token generation without KV cache, stage transitions from `</think>` to `<recaption>`, and generated COT injection via tokenizer `batch_cot_text` for `gen_image`.
+
+**Tech Stack:** Python, PyTorch, Hunyuan tokenizer/image processor helpers, LightX2V native HunyuanImage3 model/infer/weights.
+
+---
+
+### Task 1: Regression Tests
+
+**Files:**
+- Modify: `tests/test_hunyuan_image3_integration.py`
+
+- [ ] Add tests for resolving bot task, generating COT before image inputs, and stage transition token forcing.
+- [ ] Run targeted pytest and verify failures before implementation.
+
+### Task 2: Runner Text Generation Helpers
+
+**Files:**
+- Modify: `lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py`
+
+- [ ] Add helpers for system prompt resolution, text template creation, next-token selection, and COT text generation.
+- [ ] Preserve existing `bot_task=image` behavior.
+- [ ] For `think_recaption`, force `<recaption>` after `</think>` and stop after `</recaption>`.
+
+### Task 3: Feed COT Into Image Generation
+
+**Files:**
+- Modify: `lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py`
+- Modify: `configs/hunyuan_image3/hunyuan_image3_t2i.json`
+
+- [ ] Pass `batch_cot_text` into `apply_chat_template(..., mode="gen_image")`.
+- [ ] Add config keys for `bot_task`, `use_system_prompt`, `max_new_tokens`, and deterministic text generation.
+
+### Task 4: Verification
+
+**Files:**
+- Test: `tests/test_hunyuan_image3_integration.py`
+
+- [ ] Run full Hunyuan integration tests.
+- [ ] Run ruff and py_compile on touched files.

lightx2v/common/ops/norm/rms_norm_weight.py

@@ -66,14 +66,17 @@ def _get_lora_attr_mapping(self):
         self.weight_diff = torch.tensor(0.0, dtype=GET_DTYPE(), device=AI_DEVICE)
 
     def _get_actual_weight(self):
-        if not hasattr(self, "weight_diff"):
+        if not hasattr(self, "weight_diff") or not self.has_diff:
             return self.weight
+        if self.weight_diff.device != self.weight.device or self.weight_diff.dtype != self.weight.dtype:
+            self.weight_diff = self.weight_diff.to(device=self.weight.device, dtype=self.weight.dtype)
         return self.weight + self.weight_diff
 
     def register_diff(self, weight_dict):
         if not self.lazy_load or self.create_cuda_buffer or self.create_cpu_buffer:
             if self.weight_diff_name in weight_dict:
                 self.weight_diff = weight_dict[self.weight_diff_name]
+                self.has_diff = True
                 logger.debug(f"Register Diff to {self.weight_name}")
 
     def load(self, weight_dict):

lightx2v/infer.py

@@ -12,6 +12,7 @@
 from lightx2v.models.runners.flux2.flux2_runner import Flux2DevRunner, Flux2KleinRunner  # noqa: F401
 from lightx2v.models.runners.hidream_o1_image.hidream_o1_image_runner import HidreamO1ImageRunner  # noqa: F401
 from lightx2v.models.runners.hunyuan3d.hunyuan3d_shape_runner import Hunyuan3DShapeRunner  # noqa: F401
+from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner  # noqa: F401
 from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_distill_runner import HunyuanVideo15DistillRunner  # noqa: F401
 from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner  # noqa: F401
 from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner  # noqa: F401
@@ -47,6 +48,17 @@
 from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER
 
 
+def _str2bool(value):
+    if isinstance(value, bool):
+        return value
+    lowered = value.lower()
+    if lowered in ("1", "true", "yes", "y", "on"):
+        return True
+    if lowered in ("0", "false", "no", "n", "off"):
+        return False
+    raise argparse.ArgumentTypeError(f"Expected a boolean value, got {value!r}")
+
+
 def init_runner(config):
     torch.set_grad_enabled(False)
     runner = RUNNER_REGISTER[config["model_cls"]](config)
@@ -88,6 +100,7 @@ def main():
             "wan2.2_s2v",
             "hunyuan_video_1.5",
             "hunyuan_video_1.5_distill",
+            "hunyuan_image3",
             "hunyuan3d",
             "worldplay_distill",
             "worldplay_ar",
@@ -116,6 +129,36 @@ def main():
     parser.add_argument("--support_tasks", type=str, nargs="+", default=[], help="Set supported tasks for the model")
     parser.add_argument("--model_path", type=str, required=True)
     parser.add_argument("--config_json", type=str, required=True)
+    hunyuan_cache_args = parser.add_argument_group("HunyuanImage3 cache options")
+    hunyuan_cache_args.add_argument("--enable_kv_cache", "--enable-kv-cache", dest="enable_kv_cache", type=_str2bool, nargs="?", const=True, default=None)
+    hunyuan_cache_args.add_argument("--enable_text_kv_cache", "--enable-text-kv-cache", dest="enable_text_kv_cache", type=_str2bool, nargs="?", const=True, default=None)
+    hunyuan_cache_args.add_argument("--use_taylor_cache", "--use-taylor-cache", dest="use_taylor_cache", type=_str2bool, nargs="?", const=True, default=None)
+    hunyuan_cache_args.add_argument("--taylor_cache_interval", "--taylor-cache-interval", dest="taylor_cache_interval", type=int, default=None)
+    hunyuan_cache_args.add_argument("--taylor_cache_order", "--taylor-cache-order", dest="taylor_cache_order", type=int, default=None)
+    hunyuan_cache_args.add_argument(
+        "--taylor_cache_enable_first_enhance",
+        "--taylor-cache-enable-first-enhance",
+        dest="taylor_cache_enable_first_enhance",
+        type=_str2bool,
+        nargs="?",
+        const=True,
+        default=None,
+    )
+    hunyuan_cache_args.add_argument("--taylor_cache_first_enhance_steps", "--taylor-cache-first-enhance-steps", dest="taylor_cache_first_enhance_steps", type=int, default=None)
+    hunyuan_cache_args.add_argument(
+        "--taylor_cache_enable_tailing_enhance",
+        "--taylor-cache-enable-tailing-enhance",
+        dest="taylor_cache_enable_tailing_enhance",
+        type=_str2bool,
+        nargs="?",
+        const=True,
+        default=None,
+    )
+    hunyuan_cache_args.add_argument("--taylor_cache_tailing_enhance_steps", "--taylor-cache-tailing-enhance-steps", dest="taylor_cache_tailing_enhance_steps", type=int, default=None)
+    hunyuan_cache_args.add_argument("--taylor_cache_low_freqs_order", "--taylor-cache-low-freqs-order", dest="taylor_cache_low_freqs_order", type=int, default=None)
+    hunyuan_cache_args.add_argument("--taylor_cache_high_freqs_order", "--taylor-cache-high-freqs-order", dest="taylor_cache_high_freqs_order", type=int, default=None)
+    hunyuan_native_args = parser.add_argument_group("HunyuanImage3 native implementation options")
+    hunyuan_native_args.add_argument("--moe_impl", "--moe-impl", dest="moe_impl", type=str, choices=["eager", "flashinfer"], default=None)
     parser.add_argument("--use_prompt_enhancer", action="store_true")
     parser.add_argument("--prompt", type=str, default="", help="The input prompt for text-to-video generation")
     parser.add_argument("--negative_prompt", type=str, default="")

lightx2v/models/networks/hunyuan_image3/__init__.py

@@ -0,0 +1,3 @@
+from .model import HunyuanImage3Model
+
+__all__ = ["HunyuanImage3Model"]

lightx2v/models/networks/hunyuan_image3/infer/kv_cache.py

@@ -0,0 +1,130 @@
+from dataclasses import dataclass
+import math
+
+import torch
+
+
+@dataclass
+class HunyuanImage3KVCacheLayer:
+    key: torch.Tensor | None = None
+    value: torch.Tensor | None = None
+
+
+class HunyuanImage3StaticKVCache:
+    """Per-layer KV cache matching HunyuanImage3 gen_text/gen_image inference."""
+
+    def __init__(self, num_layers, max_cache_len, dynamic=False):
+        self.num_layers = int(num_layers)
+        self.max_cache_len = int(max_cache_len)
+        self.dynamic = bool(dynamic)
+        self.layers = [HunyuanImage3KVCacheLayer() for _ in range(self.num_layers)]
+
+    def _ensure_layer(self, layer_idx, key_states, value_states):
+        layer = self.layers[layer_idx]
+        if layer.key is None:
+            key_shape = (*key_states.shape[:2], self.max_cache_len, key_states.shape[-1])
+            value_shape = (*value_states.shape[:2], self.max_cache_len, value_states.shape[-1])
+            layer.key = torch.zeros(key_shape, device=key_states.device, dtype=key_states.dtype)
+            layer.value = torch.zeros(value_shape, device=value_states.device, dtype=value_states.dtype)
+        return layer
+
+    def update(self, key_states, value_states, layer_idx, cache_position=None):
+        layer = self._ensure_layer(layer_idx, key_states, value_states)
+        if cache_position is None:
+            layer.key[:, :, : key_states.shape[2]].copy_(key_states)
+            layer.value[:, :, : value_states.shape[2]].copy_(value_states)
+            return self._slice_dynamic(layer, key_states.shape[2])
+
+        cache_position = cache_position.to(device=key_states.device, dtype=torch.long)
+        if cache_position.dim() == 1:
+            layer.key.index_copy_(2, cache_position, key_states)
+            layer.value.index_copy_(2, cache_position, value_states)
+            return self._slice_dynamic(layer, int(cache_position[-1].item()) + 1)
+
+        if cache_position.dim() != 2:
+            raise ValueError(f"HunyuanImage3 cache_position must be 1D or 2D, got {cache_position.shape}.")
+        if cache_position.shape[0] != key_states.shape[0]:
+            raise ValueError(
+                f"HunyuanImage3 cache batch mismatch: cache_position={cache_position.shape}, key_states={key_states.shape}."
+            )
+
+        for batch_idx in range(cache_position.shape[0]):
+            layer.key[batch_idx].index_copy_(1, cache_position[batch_idx], key_states[batch_idx])
+            layer.value[batch_idx].index_copy_(1, cache_position[batch_idx], value_states[batch_idx])
+        return self._slice_dynamic(layer, int(cache_position.max().item()) + 1)
+
+    def _slice_dynamic(self, layer, end):
+        if not self.dynamic:
+            return layer.key, layer.value
+        end = min(int(end), self.max_cache_len)
+        return layer.key[:, :, :end], layer.value[:, :, :end]
+
+
+def _decompose_freqs(x, cutoff_ratio=0.1):
+    original_dtype = x.dtype
+    x_fp32 = x.float()
+    freq = torch.fft.fft(x_fp32, dim=1)
+    freqs = torch.fft.fftfreq(x_fp32.shape[1], d=1.0, device=x.device)
+    cutoff = cutoff_ratio * freqs.abs().max()
+    low_mask = (freqs.abs() <= cutoff)[None, :, None]
+    high_mask = ~low_mask
+    low = torch.fft.ifft(freq * low_mask, dim=1).real.to(dtype=original_dtype)
+    high = torch.fft.ifft(freq * high_mask, dim=1).real.to(dtype=original_dtype)
+    return low, high
+
+
+class HunyuanImage3TaylorCache:
+    """Frequency-split Taylor hidden-state cache used by HunyuanImage3 sampling."""
+
+    def __init__(self, max_order):
+        self.max_order = int(max_order)
+        self.low_derivatives = [None for _ in range(self.max_order + 1)]
+        self.high_derivatives = [None for _ in range(self.max_order + 1)]
+        self.last_past_key_values = None
+
+    def taylor_formula(self, distance):
+        low_output = None
+        high_output = None
+        for order, derivative in enumerate(self.low_derivatives):
+            if derivative is None:
+                break
+            term = (distance**order / math.factorial(order)) * derivative
+            low_output = term if low_output is None else low_output + term
+        for order, derivative in enumerate(self.high_derivatives):
+            if derivative is None:
+                break
+            term = (distance**order / math.factorial(order)) * derivative
+            high_output = term if high_output is None else high_output + term
+        if low_output is None and high_output is None:
+            raise RuntimeError("HunyuanImage3 Taylor cache has no derivatives to extrapolate from.")
+        if low_output is None:
+            return high_output
+        if high_output is None:
+            return low_output
+        return low_output + high_output
+
+    def derivatives_computation(self, hidden_states, distance, low_freqs_order, high_freqs_order):
+        low, high = _decompose_freqs(hidden_states)
+        new_low = [None for _ in range(self.max_order + 1)]
+        new_high = [None for _ in range(self.max_order + 1)]
+        new_low[0] = low
+        new_high[0] = high
+        safe_distance = max(int(distance), 1)
+
+        for order in range(min(int(low_freqs_order), self.max_order)):
+            if self.low_derivatives[order] is None:
+                break
+            new_low[order + 1] = (new_low[order] - self.low_derivatives[order]) / safe_distance
+
+        for order in range(min(int(high_freqs_order), self.max_order)):
+            if self.high_derivatives[order] is None:
+                break
+            new_high[order + 1] = (new_high[order] - self.high_derivatives[order]) / safe_distance
+
+        self.low_derivatives = new_low
+        self.high_derivatives = new_high
+
+    def clear_derivatives(self):
+        self.low_derivatives = [None for _ in range(self.max_order + 1)]
+        self.high_derivatives = [None for _ in range(self.max_order + 1)]
+        self.last_past_key_values = None

lightx2v/models/networks/hunyuan_image3/infer/module_io.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class HunyuanImage3PreInferOutput:
+    hidden_states: torch.Tensor
+    attention_mask: torch.Tensor | None = None
+    position_ids: torch.Tensor | None = None
+    custom_pos_emb: tuple[torch.Tensor, torch.Tensor] | None = None
+    past_key_values: object | None = None
+    use_cache: bool = False
+    image_mask: torch.Tensor | None = None
+    timesteps: torch.Tensor | None = None
+    token_hw: tuple[int, int] | None = None
+    first_step: bool | None = None