from dataclasses import dataclass
from types import MethodType
from time import sleep

import torch
from torch import nn
from torch import Tensor
from einops import rearrange

from comfy.ldm.flux.model import Flux
from comfy.ldm.hunyuan_video.model import HunyuanVideo
from comfy.ldm.flux.layers import DoubleStreamBlock, SingleStreamBlock, LastLayer
from comfy.ldm.wan.model import (
    WanModel,
    VaceWanModel,
    CameraWanModel,
    WanModel_S2V,
    HumoWanModel,
    WanAttentionBlock,
    VaceWanAttentionBlock,
    Head,
)
from comfy.ldm.flux.math import apply_rope1
from comfy.ldm.modules.attention import optimized_attention
from comfy.model_patcher import ModelPatcher
import comfy.model_management


def apply_pe(x: Tensor, pe: Tensor) -> Tensor:
    """
    The PE application from flux.math.attention, removed, so that we can cache the keys post-PE
    """
    shape = x.shape
    dtype = x.dtype
    x = x.float().reshape(*x.shape[:-1], -1, 1, 2)
    x = (pe[..., 0] * x[..., 0] + pe[..., 1] * x[..., 1]).reshape(*shape).to(dtype)

    return x


def take_attributes_from(source, target, keys):
    for x in keys:
        setattr(target, x, getattr(source, x))


@dataclass
class RASConfig:
    warmup_steps: int = 4
    hydrate_every: int = 5
    sample_ratio: float = 0.5
    starvation_scale: float = 0.1
    high_ratio: float = 1.0


class RASManager:
    """
    Coordinates the live indices, metrics, and model wrapping.
    """

    def __init__(self, config: RASConfig):
        self.flipped_img_txt = False
        self.timestep: int = 0
        self.n_txt: int = 0
        self.n_img: int = 0
        self.cached_output: Tensor | None = None
        self.live_txt_indices: Tensor | None = None
        self.live_img_indices: Tensor | None = None
        self.drop_count: torch.Tensor | None = None
        self.config = config
        self.patch_size: list[int]
        self.model: Flux | HunyuanVideo
        assert (
            self.config.high_ratio >= 0 and self.config.high_ratio <= 1
        ), "High ratio should be in the range of [0, 1]"

    def wrap_layer(self, layer, first: bool = False, last: bool = False):
        if isinstance(
            layer,
            (
                DoubleStreamBlockWrapper,
                SingleStreamBlockWrapper,
                LastLayerWrapper,
                WanAttentionBlockWrapper,
                VaceWanAttentionBlockWrapper,
            ),
        ):
            raise TypeError("Old wrapping wasn't removed!")

        if isinstance(layer, DoubleStreamBlock):
            wrapped = DoubleStreamBlockWrapper(layer, self, first)
        elif isinstance(layer, SingleStreamBlock):
            wrapped = SingleStreamBlockWrapper(layer, self, last)
        elif isinstance(layer, LastLayer):
            wrapped = LastLayerWrapper(layer, self)
        elif isinstance(layer, WanAttentionBlock):
            wrapped = WanAttentionBlockWrapper(layer, self, first, last)
        elif isinstance(layer, VaceWanAttentionBlock):
            wrapped = VaceWanAttentionBlockWrapper(layer, self, first, last)
        elif isinstance(layer, Head):
            wrapped = HeadWrapper(layer, self)
        else:
            raise TypeError(f"Can't wrap layer of type {layer.__class__.__name__}")
        return wrapped

    def wrap_model(self, patcher: ModelPatcher):
        model = patcher.model.diffusion_model
        self.model = model
        if isinstance(model, Flux):
            self.patch_size = [model.patch_size, model.patch_size]
        elif isinstance(model, HunyuanVideo):
            self.patch_size = model.patch_size
        elif isinstance(
            model, (WanModel, VaceWanModel, CameraWanModel, WanModel_S2V, HumoWanModel)
        ):
            self.patch_size = model.patch_size
        else:
            raise TypeError(f"Can't wrap model of type {model.__class__.__name__}")
        # Handle different model architectures
        if isinstance(model, (Flux, HunyuanVideo)):
            # wrap the single and double blocks to have caching
            for i, v in enumerate(model.double_blocks):
                # first block has the special responsibility of removing tokens
                if i == 0:
                    self.flipped_img_txt = v.flipped_img_txt
                    layer = self.wrap_layer(v, first=True)
                else:
                    layer = self.wrap_layer(v)
                patcher.add_object_patch(f"diffusion_model.double_blocks.{i}", layer)
            for i, v in enumerate(model.single_blocks):
                # last block will put them back
                patcher.add_object_patch(
                    f"diffusion_model.single_blocks.{i}",
                    self.wrap_layer(v, last=(i == (len(model.single_blocks) - 1))),
                )
        elif isinstance(
            model, (WanModel, VaceWanModel, CameraWanModel, WanModel_S2V, HumoWanModel)
        ):
            # Wan models have a different structure with just 'blocks'
            self.flipped_img_txt = False  # Wan doesn't use the flipped pattern
            for i, v in enumerate(model.blocks):
                # first block has the special responsibility of removing tokens
                # last block will put them back
                layer = self.wrap_layer(
                    v, first=(i == 0), last=(i == (len(model.blocks) - 1))
                )
                patcher.add_object_patch(f"diffusion_model.blocks.{i}", layer)

        # wrap the forward_orig method, to be able to get the timestep
        forward_orig = model.forward_orig

        def new_forward(_self, *args, **kwargs):
            # Get transformer_options from kwargs for Wan models, or from args for Flux/Hunyuan
            if "transformer_options" in kwargs:
                transformer_options = kwargs["transformer_options"]
            else:
                transformer_options = args[-1]

            self.timestep = self.timestep_from_sigmas(
                transformer_options["sigmas"], transformer_options["sample_sigmas"]
            )
            if self.timestep == 0:
                # reset as much as possible
                self.live_img_indices = None
                self.live_txt_indices = None
                self.drop_count = None
            return forward_orig(*args, **kwargs)

        patcher.add_object_patch(
            "diffusion_model.forward_orig", MethodType(new_forward, model)
        )
        # wrap the last_layer, to be able to read the output and calculate the metric
        if isinstance(model, (Flux, HunyuanVideo)):
            patcher.add_object_patch(
                "diffusion_model.final_layer", self.wrap_layer(model.final_layer)
            )
        elif isinstance(
            model, (WanModel, VaceWanModel, CameraWanModel, WanModel_S2V, HumoWanModel)
        ):
            # Wan models use 'head' instead of 'final_layer'
            patcher.add_object_patch(
                "diffusion_model.head", self.wrap_layer(model.head)
            )

    @staticmethod
    def timestep_from_sigmas(sigmas: Tensor, sample_sigmas: Tensor):
        # we assume that one element of sample_sigmas is exactly equal to sigmas
        # but we'll still check explicitly, using an argmin, in case of some loss of precision
        s = sigmas.item()
        i = torch.argmin(torch.abs(sample_sigmas - s).flatten())
        return int(i.item())

    def skip_ratio(self, timestep: int) -> float:
        if timestep < self.config.warmup_steps:
            return 0

        if self.config.hydrate_every:
            if (
                1 + timestep - self.config.warmup_steps
            ) % self.config.hydrate_every == 0:
                return 0
        result = 1.0 - self.config.sample_ratio
        return result

    def select_indices(self, diff: Tensor, timestep: int):
        if isinstance(self.model, Flux):
            # b (h w) (c ph pw) = model_out.shape
            metric = rearrange(
                diff,
                "b s (c ph pw) -> b s ph pw c",
                ph=self.patch_size[0],
                pw=self.patch_size[1],
            )
            metric = torch.std(metric, dim=-1).mean((-1, -2))
        elif isinstance(
            self.model,
            (
                HunyuanVideo,
                WanModel,
                VaceWanModel,
                CameraWanModel,
                WanModel_S2V,
                HumoWanModel,
            ),
        ):
            # b (h w) (c ph pw) = model_out.shape
            # Both HunyuanVideo and Wan models are video models with 3D patches
            metric = rearrange(
                diff,
                "b s (c pt ph pw ) -> b s pt ph pw c",
                pt=self.patch_size[0],
                ph=self.patch_size[1],
                pw=self.patch_size[2],
            )
            metric = torch.std(metric, dim=-1).mean((-1, -2, -3))
        else:
            raise TypeError("Unknown latent type!")
        # for batch size > 1, we pick separate indices per batch
        # for now, JUST FOR TESTING, we'll merge all the batches and use the indices that are the most relevant for all batches
        metric = metric.mean(dim=0)
        metric = metric.flatten()
        if self.drop_count is None:
            self.drop_count = torch.zeros(
                metric.shape, dtype=torch.int, device=diff.device
            )
        # hmm, what if we do a gaussian blur or some sort of spatial lowpass, to improve the spatial continuity of the patches?
        metric *= torch.exp(self.config.starvation_scale * self.drop_count)
        indices = torch.sort(metric, dim=-1, descending=False).indices
        skip_ratio = self.skip_ratio(timestep)
        if skip_ratio <= 0.01:
            # we're not dropping anything -- remove the live_indices
            # we use the value None to indicate a full hydrate
            self.live_img_indices = None
        else:
            low_bar = int(skip_ratio * len(metric) * (1 - self.config.high_ratio))
            high_bar = int(skip_ratio * len(metric) * self.config.high_ratio)
            cache_indices = torch.cat([indices[:low_bar], indices[-high_bar:]])
            self.live_img_indices = indices[low_bar:-high_bar]
            self.drop_count[cache_indices] += 1

        # TODO: for now we keep all txt tokens
        # in the future, we can probably do something like randomly keep a fraction of them
        if self.n_txt > 0:
            self.live_txt_indices = torch.arange(
                0, self.n_txt, dtype=torch.int, device=diff.device
            )

    def live_indices(self):
        if self.live_img_indices is None or self.live_txt_indices is None:
            return self.live_img_indices
        if self.flipped_img_txt:
            result = torch.cat(
                (self.live_img_indices, self.live_txt_indices + self.n_img)
            )
        else:
            result = torch.cat(
                (self.live_txt_indices, self.live_img_indices + self.n_txt)
            )
        return result


class DoubleStreamBlockWrapper(DoubleStreamBlock):
    """
    Same as the DoubleStreamBlock, but uses a RASManager and RASCache to do KV caching.
    """

    def __init__(self, original: DoubleStreamBlock, manager: RASManager, first=False):
        nn.Module.__init__(self)
        take_attributes_from(
            original,
            self,
            [
                "num_heads",
                "hidden_size",
                "img_mod",
                "img_norm1",
                "img_attn",
                "img_norm2",
                "img_mlp",
                "txt_mod",
                "txt_norm1",
                "txt_attn",
                "txt_norm2",
                "txt_mlp",
                "flipped_img_txt",
            ],
        )
        self.manager = manager
        self.k_cache: torch.Tensor
        self.v_cache: torch.Tensor
        self.first = first

    def forward(self, img, txt, vec, pe, attn_mask=None):
        # RAS: if this is the first doublestreamblock, then we should drop some of the img and txt tokens
        idx = self.manager.live_indices()
        if self.first:
            self.manager.n_txt = txt.shape[1]
            self.manager.n_img = img.shape[1]

            img_idx = self.manager.live_img_indices
            txt_idx = self.manager.live_txt_indices
            if idx is not None:
                img = img[..., img_idx, :]
                txt = txt[..., txt_idx, :]

        img_mod1, img_mod2 = self.img_mod(vec)
        txt_mod1, txt_mod2 = self.txt_mod(vec)

        # prepare image for attention
        img_modulated = self.img_norm1(img)
        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
        img_qkv = self.img_attn.qkv(img_modulated)
        img_q, img_k, img_v = img_qkv.view(
            img_qkv.shape[0], img_qkv.shape[1], 3, self.num_heads, -1
        ).permute(2, 0, 3, 1, 4)
        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)

        # prepare txt for attention
        txt_modulated = self.txt_norm1(txt)
        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
        txt_qkv = self.txt_attn.qkv(txt_modulated)
        txt_q, txt_k, txt_v = txt_qkv.view(
            txt_qkv.shape[0], txt_qkv.shape[1], 3, self.num_heads, -1
        ).permute(2, 0, 3, 1, 4)
        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)

        # RAS: KV Cache and Attention Call
        # select part of the PE
        if idx is not None:
            pe = pe[:, :, idx]

        # create queries, keys, and values
        if self.flipped_img_txt:
            queries = apply_pe(torch.cat((img_q, txt_q), dim=2), pe)
            keys = apply_pe(torch.cat((img_k, txt_k), dim=2), pe)
            values = torch.cat((img_v, txt_v), dim=2)
        else:
            queries = apply_pe(torch.cat((txt_q, img_q), dim=2), pe)
            keys = apply_pe(torch.cat((txt_k, img_k), dim=2), pe)
            values = torch.cat((txt_v, img_v), dim=2)

        # fill in the KV cache
        if idx is None:
            self.k_cache = keys
            self.v_cache = values
        else:
            self.k_cache[..., idx, :] = keys
            self.v_cache[..., idx, :] = values
        # actual attention call
        attn = optimized_attention(
            queries,
            self.k_cache,
            self.v_cache,
            img_q.shape[1],
            skip_reshape=True,
            mask=attn_mask,
        )
        if self.flipped_img_txt:
            img_attn, txt_attn = attn[:, : img.shape[1]], attn[:, img.shape[1] :]
        else:
            txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
        # End of RAS code

        # calculate the img blocks
        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
        img = img + img_mod2.gate * self.img_mlp(
            (1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift
        )

        # calculate the txt blocks
        txt += txt_mod1.gate * self.txt_attn.proj(txt_attn)
        txt += txt_mod2.gate * self.txt_mlp(
            (1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift
        )

        if txt.dtype == torch.float16:
            txt = torch.nan_to_num(txt, nan=0.0, posinf=65504, neginf=-65504)

        return img, txt


class SingleStreamBlockWrapper(SingleStreamBlock):
    """
    Same as the SingleStreamBlock, but uses a RASManager and RASCache to do KV caching.
    """

    def __init__(self, original: SingleStreamBlock, manager: RASManager, last=False):
        nn.Module.__init__(self)
        # steal all the attributes from the SingleStreamBlock
        take_attributes_from(
            original,
            self,
            [
                "hidden_dim",
                "num_heads",
                "scale",
                "mlp_hidden_dim",
                "linear1",
                "linear2",
                "norm",
                "hidden_size",
                "pre_norm",
                "mlp_act",
                "modulation",
            ],
        )
        self.manager = manager
        self.k_cache: torch.Tensor
        self.v_cache: torch.Tensor
        self.last = last

    def forward(self, x, vec, pe, attn_mask=None):
        idx = self.manager.live_indices()
        mod, _ = self.modulation(vec)
        qkv, mlp = torch.split(
            self.linear1((1 + mod.scale) * self.pre_norm(x) + mod.shift),
            [3 * self.hidden_size, self.mlp_hidden_dim],
            dim=-1,
        )

        q, k, v = qkv.view(qkv.shape[0], qkv.shape[1], 3, self.num_heads, -1).permute(
            2, 0, 3, 1, 4
        )
        q, k = self.norm(q, k, v)

        # RAS: KV Cache
        if idx is not None:
            pe = pe[:, :, idx]
        q = apply_pe(q, pe)
        k = apply_pe(k, pe)
        # full hydrate
        if idx is None:
            self.k_cache = k
            self.v_cache = v
        # partial update
        else:
            self.k_cache[..., idx, :] = k
            self.v_cache[..., idx, :] = v
        attn = optimized_attention(
            q,
            self.k_cache,
            self.v_cache,
            q.shape[1],
            skip_reshape=True,
            mask=attn_mask,
        )
        # End of RAS code
        # compute activation in mlp stream, cat again and run second linear layer
        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
        x += mod.gate * output
        if x.dtype == torch.float16:
            x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)

        if self.last:
            # put the relevant tokens back into the cached output
            if idx is None or self.manager.cached_output is None:
                self.manager.cached_output = x.clone()
            else:
                self.manager.cached_output[..., idx, :] = x
            return self.manager.cached_output
        return x


class LastLayerWrapper(LastLayer):
    """
    Same as the LastLayer, but reports its output to a manager.
    """

    def __init__(self, original: LastLayer, manager: RASManager):
        nn.Module.__init__(self)
        take_attributes_from(
            original, self, ["norm_final", "linear", "adaLN_modulation"]
        )
        self.manager = manager

    def forward(self, x, vec) -> Tensor:
        output = super().forward(x, vec)
        self.manager.select_indices(output, self.manager.timestep)
        return output


class WanAttentionBlockWrapper(WanAttentionBlock):
    """
    Same as the WanAttentionBlock, but uses a RASManager and RASCache to do KV caching.
    """

    def __init__(self, original: WanAttentionBlock, manager: RASManager, first, last):
        nn.Module.__init__(self)
        take_attributes_from(
            original,
            self,
            [
                "dim",
                "ffn_dim",
                "num_heads",
                "window_size",
                "qk_norm",
                "cross_attn_norm",
                "eps",
                "norm1",
                "self_attn",
                "norm3",
                "cross_attn",
                "norm2",
                "ffn",
                "modulation",
            ],
        )
        self.manager = manager
        self.k_cache: torch.Tensor
        self.v_cache: torch.Tensor
        self.first = first
        self.last = last

    def forward(
        self, x, e, freqs, context, context_img_len=257, transformer_options={}
    ):
        # RAS: if this is the first block, then we should drop some tokens
        idx = self.manager.live_indices()

        if self.first:
            # we just have img tokens
            self.manager.n_txt = 0
            self.manager.n_img = x.shape[1]
            if idx is not None:
                x = x[..., idx, :]

        # Modulation handling (copied from original)
        if e.ndim < 4:
            e = (
                comfy.model_management.cast_to(
                    self.modulation, dtype=x.dtype, device=x.device
                )
                + e
            ).chunk(6, dim=1)
        else:
            e = (
                comfy.model_management.cast_to(
                    self.modulation, dtype=x.dtype, device=x.device
                ).unsqueeze(0)
                + e
            ).unbind(2)

        # Self-attention with RAS caching
        y = self.self_attn_with_cache(
            torch.addcmul(
                self.repeat_e(e[0], x), self.norm1(x), 1 + self.repeat_e(e[1], x)
            ),
            freqs,
            idx,
            transformer_options=transformer_options,
        )

        x = torch.addcmul(x, y, self.repeat_e(e[2], x))

        # Cross-attention & ffn (unchanged from original)
        x = x + self.cross_attn(
            self.norm3(x),
            context,
            context_img_len=context_img_len,
            transformer_options=transformer_options,
        )
        y = self.ffn(
            torch.addcmul(
                self.repeat_e(e[3], x), self.norm2(x), 1 + self.repeat_e(e[4], x)
            )
        )
        x = torch.addcmul(x, y, self.repeat_e(e[5], x))

        if self.last:
            # put the relevant tokens back into the cached output
            if idx is None or self.manager.cached_output is None:
                self.manager.cached_output = x.clone()
            else:
                self.manager.cached_output[..., idx, :] = x
            return self.manager.cached_output
        return x

    def repeat_e(self, e, x):
        """Helper function for modulation broadcasting"""
        repeats = 1
        if e.size(1) > 1:
            repeats = x.size(1) // e.size(1)
        if repeats == 1:
            return e
        if repeats * e.size(1) == x.size(1):
            return torch.repeat_interleave(e, repeats, dim=1)
        else:
            return torch.repeat_interleave(e, repeats + 1, dim=1)[:, : x.size(1)]

    def self_attn_with_cache(self, x, freqs, idx, transformer_options={}):
        """Modified self-attention that uses KV caching"""
        b, s, n, d = *x.shape[:2], self.num_heads, self.self_attn.head_dim

        # just pull out part of the frequencies
        if idx is not None:
            freqs = freqs[:, idx]

        # Compute QKV like original Wan self-attention
        q = self.self_attn.norm_q(self.self_attn.q(x)).view(b, s, n, d)
        q = apply_rope1(q, freqs).view(b, s, n * d)
        k = self.self_attn.norm_k(self.self_attn.k(x)).view(b, s, n, d)
        k = apply_rope1(k, freqs).view(b, s, n * d)
        v = self.self_attn.v(x).view(b, s, n * d)

        # RAS: KV Cache management
        if idx is None:
            self.k_cache = k
            self.v_cache = v
        else:
            self.k_cache[:, idx, :] = k
            self.v_cache[:, idx, :] = v
        x = optimized_attention(
            q,
            self.k_cache,
            self.v_cache,
            heads=n,
            transformer_options=transformer_options,
        )

        x = self.self_attn.o(x)
        return x


class VaceWanAttentionBlockWrapper(VaceWanAttentionBlock):
    """
    Same as the VaceWanAttentionBlock, but uses a RASManager and RASCache to do KV caching.
    """

    def __init__(
        self, original: VaceWanAttentionBlock, manager: RASManager, first=False
    ):
        nn.Module.__init__(self)
        # Copy attributes, handling the case where some might not exist
        attributes_to_copy = [
            "dim",
            "ffn_dim",
            "num_heads",
            "window_size",
            "qk_norm",
            "cross_attn_norm",
            "eps",
            "block_id",
            "norm1",
            "self_attn",
            "norm3",
            "cross_attn",
            "norm2",
            "ffn",
            "modulation",
        ]

        # Handle optional attributes that might not exist
        optional_attributes = ["before_proj", "after_proj"]

        take_attributes_from(original, self, attributes_to_copy)

        # Copy optional attributes if they exist
        for attr in optional_attributes:
            if hasattr(original, attr):
                setattr(self, attr, getattr(original, attr))
        self.manager = manager
        self.k_cache: torch.Tensor
        self.v_cache: torch.Tensor
        self.first = first

    def forward(
        self, c, x, e, freqs, context, context_img_len=257, transformer_options={}
    ):
        # RAS: if this is the first block, then we should drop some tokens
        idx = self.manager.live_indices()
        if self.first:
            # For VaceWan models, we track the sequence length differently
            self.manager.n_txt = (
                0  # Wan doesn't have separate text tokens in the same way
            )
            # TODO: get vace working
            self.manager.n_img = c.shape[1]

            if idx is not None:
                c = c[..., idx, :]

        # VaceWan specific logic (copied from original)
        if self.block_id == 0 and hasattr(self, "before_proj"):
            c = self.before_proj(c) + x

        # Call the parent WanAttentionBlock logic but with our modified attention
        # Modulation handling (copied from original)
        if e.ndim < 4:
            e = (
                comfy.model_management.cast_to(
                    self.modulation, dtype=c.dtype, device=c.device
                )
                + e
            ).chunk(6, dim=1)
        else:
            e = (
                comfy.model_management.cast_to(
                    self.modulation, dtype=c.dtype, device=c.device
                ).unsqueeze(0)
                + e
            ).unbind(2)

        # Self-attention with RAS caching
        y = self.self_attn_with_cache(
            torch.addcmul(
                self.repeat_e(e[0], c), self.norm1(c), 1 + self.repeat_e(e[1], c)
            ),
            freqs,
            idx,
            transformer_options=transformer_options,
        )

        c = torch.addcmul(c, y, self.repeat_e(e[2], c))

        # Cross-attention & ffn (unchanged from original)
        c = c + self.cross_attn(
            self.norm3(c),
            context,
            context_img_len=context_img_len,
            transformer_options=transformer_options,
        )
        y = self.ffn(
            torch.addcmul(
                self.repeat_e(e[3], c), self.norm2(c), 1 + self.repeat_e(e[4], c)
            )
        )
        c = torch.addcmul(c, y, self.repeat_e(e[5], c))

        # VaceWan specific projection
        c_skip = self.after_proj(c)
        return c_skip, c

    def repeat_e(self, e, x):
        """Helper function for modulation broadcasting"""
        repeats = 1
        if e.size(1) > 1:
            repeats = x.size(1) // e.size(1)
        if repeats == 1:
            return e
        if repeats * e.size(1) == x.size(1):
            return torch.repeat_interleave(e, repeats, dim=1)
        else:
            return torch.repeat_interleave(e, repeats + 1, dim=1)[:, : x.size(1)]

    def self_attn_with_cache(self, x, freqs, idx, transformer_options={}):
        """Modified self-attention that uses KV caching"""
        # NOTE the following attention function doesn't work
        # it's fucked
        # needs to be replaced
        # just use the one from wan (above)
        b, s, n, d = *x.shape[:2], self.num_heads, x.shape[-1] // self.num_heads

        # Compute QKV like original Wan self-attention
        q = self.self_attn.norm_q(self.self_attn.q(x)).view(b, s, n, d)
        k = self.self_attn.norm_k(self.self_attn.k(x)).view(b, s, n, d)
        v = self.self_attn.v(x).view(b, s, n, d)

        # Apply rope to Q and K
        from comfy.ldm.flux.math import apply_rope1

        if idx is not None:
            freqs = freqs[:, idx]
        q = apply_rope1(q, freqs)
        k = apply_rope1(k, freqs)

        # RAS: KV Cache management
        if idx is None:
            # Full hydrate - store the complete k,v tensors as cache
            self.k_cache = k
            self.v_cache = v
        else:
            self.k_cache[:, idx, :, :] = k
            self.v_cache[:, idx, :, :] = v

        q_flat = q.view(b, s, n * d)
        k_flat = self.k_cache.view(b, self.k_cache.shape[1], n * d)
        v_flat = self.v_cache.view(b, self.v_cache.shape[1], n * d)

        attn = optimized_attention(
            q_flat,
            k_flat,
            v_flat,
            heads=n,
            transformer_options=transformer_options,
        )

        result = self.self_attn.o(attn)
        return result


class HeadWrapper(Head):
    """
    Same as the Head (Wan final layer), but reports its output to a manager.
    """

    def __init__(self, original: Head, manager: RASManager):
        nn.Module.__init__(self)
        take_attributes_from(
            original,
            self,
            ["dim", "out_dim", "patch_size", "eps", "norm", "head", "modulation"],
        )
        self.manager = manager

    def forward(self, x, e) -> Tensor:
        output = super().forward(x, e)
        self.manager.select_indices(output, self.manager.timestep)
        return output