#!/usr/bin/env python3
"""Block-level group offloading for a fused MiniMax-H3 transformer.

Fused blocks read their packed weight groups directly, so they stream as whole blocks
(``offload_type="block_level"``). Two parts of the model sit outside ``transformer_blocks``:

* the token refiner (two transformer blocks over the text rows, about 1.5 GB) depends only on
  the prompt embeddings. It is computed once per prompt from its own synchronously offloaded
  blocks and never occupies the GPU during the denoising forwards;
* the remaining top-level modules (embedders, output norm and projections) form the group that
  Diffusers copies to the GPU and back synchronously around every forward. The forward needs
  them on the GPU anyway, so they stay there for the whole denoising loop.
"""

from __future__ import annotations

import weakref

import torch
from diffusers.hooks import apply_group_offloading


class _LastInputCache:
    """``compute(tensor)`` for the last input tensor object, while its in-place version holds."""

    def __init__(self, compute):
        self.compute = compute
        self.key = None
        self.value = None

    def __call__(self, tensor):
        if (
            self.key is not None
            and self.key[0]() is tensor
            and self.key[1] == tensor._version
        ):
            return self.value
        self.value = self.compute(tensor)
        self.key = (weakref.ref(tensor), tensor._version)
        return self.value


def enable_fused_group_offload(transformer, **group_offload_kwargs) -> None:
    refiner = transformer.token_refiner
    del transformer.token_refiner
    transformer.enable_group_offload(**group_offload_kwargs)
    apply_group_offloading(
        refiner,
        onload_device=group_offload_kwargs["onload_device"],
        offload_device=group_offload_kwargs["offload_device"],
        offload_type="block_level",
        num_blocks_per_group=1,
    )
    transformer.token_refiner = refiner
    # The denoiser passes the same prompt-embedding tensor at every step; the context embedder
    # then returns the same tensor object, which the refiner cache recognizes.
    transformer.context_embedder.forward = _LastInputCache(
        transformer.context_embedder.forward
    )
    refiner.forward = _LastInputCache(refiner.forward)


def keep_unmatched_group_on_cuda_while_denoising(loop_wrapper_cls, transformer) -> None:
    group = transformer._diffusers_hook.get_hook("group_offloading").group
    loop_call = loop_wrapper_cls.__call__

    def call(self, components, state):
        group.offload_ = lambda: None
        try:
            return loop_call(self, components, state)
        finally:
            del group.offload_
            group.offload_()
            torch.cuda.empty_cache()

    loop_wrapper_cls.__call__ = call