"""Device-only offload workaround for the pinned bitsandbytes 0.50.2 stack.

Parent Module.to() recursively invokes _apply(), not each child's to().
Linear8bitLt.to() moves auxiliary tensors, but this path is bypassed by a
parent move. Keep CB aliases and SCB/outlier state colocated with the weight.
This is instance-local; checkpoint tensors and installed packages are unchanged.
"""
from types import MethodType
import inspect
from weakref import WeakSet

import torch


def check_sampling_support(config):
    """Fail before loading weights when an old runtime would ignore a saved grid."""
    if config.get('sample_sigmas') is not None:
        from diffusers import QwenImage21Pipeline
        if 'sample_sigmas' not in inspect.signature(QwenImage21Pipeline.__init__).parameters:
            raise RuntimeError('This checkpoint requires Diffusers sampling-grid support (PR #14950)')


def release_prefix_cache_before_decode(pipe):
    """Release finished denoising prefix tensors before the full-resolution VAE.

    The pinned upstream pipeline holds its local KV containers through decode.
    Their tensors are no longer read after the final transformer call. Clearing
    them at decode entry preserves the denoising computation and latent input.
    Weak references avoid retaining caches on latent-only calls or failures.
    """
    if getattr(pipe, '_prefix_cache_release', False):
        return
    caches = WeakSet()
    telemetry = {'calls': 0, 'tensor_bytes': 0}

    def capture(_module, _args, kwargs):
        cache = kwargs.get('kv_cache')
        if cache is not None:
            caches.add(cache)

    handle = pipe.transformer.register_forward_pre_hook(capture, with_kwargs=True)
    original = pipe.vae.decode

    def decode(self, *args, **kwargs):
        released = 0
        for cache in list(caches):
            for layer in cache.layer_caches:
                released += sum(t.numel() * t.element_size() for t in (layer.k, layer.v)
                                if isinstance(t, torch.Tensor))
                layer.k = None
                layer.v = None
        caches.clear()
        telemetry.update(calls=telemetry['calls'] + 1, tensor_bytes=released)
        return original(*args, **kwargs)

    pipe.vae.decode = MethodType(decode, pipe.vae)
    pipe._prefix_cache_release_hook = handle
    pipe._prefix_cache_release = True
    pipe._prefix_cache_telemetry = telemetry


def patch_int8_device_moves(model):
    from bitsandbytes.nn import Linear8bitLt
    for layer in model.modules():
        if not isinstance(layer, Linear8bitLt) or getattr(layer, '_int8_move_patched', False):
            continue
        if layer.state.has_fp16_weights:
            raise ValueError('Offload helper only supports inference INT8 weights')
        original = layer._apply

        def apply(self, fn, recurse=True, _original=original):
            result = _original(fn, recurse=recurse)
            device = self.weight.device
            # Parameter auxiliaries are used before init_8bit_state(); state
            # auxiliaries are used afterwards. Never duplicate the CB storage.
            if self.weight.CB is not None:
                self.weight.CB = self.weight.data
            if self.weight.SCB is not None:
                self.weight.SCB = self.weight.SCB.to(device)
            for name, value in vars(self.state).items():
                if isinstance(value, torch.Tensor):
                    setattr(self.state, name, self.weight.data if name == 'CB' else value.to(device))
            return result

        layer._apply = MethodType(apply, layer)
        layer._int8_move_patched = True


def enable_int8_cpu_offload(pipe):
    """Install auxiliary-tensor movement before standard model CPU offload."""
    for name in ('text_encoder', 'transformer'):
        patch_int8_device_moves(getattr(pipe, name))
    pipe.enable_model_cpu_offload()


def load_int8_pipeline(model, local_files_only=False):
    """Load/offload components sequentially to avoid a combined CUDA load peak."""
    from diffusers import QwenImage21Pipeline, QwenImage21Transformer2DModel
    from transformers import Qwen3VLForConditionalGeneration
    check_sampling_support(QwenImage21Pipeline.load_config(model, local_files_only=local_files_only))
    components = {}
    for name, cls in [('transformer', QwenImage21Transformer2DModel),
                      ('text_encoder', Qwen3VLForConditionalGeneration)]:
        component = cls.from_pretrained(model, subfolder=name, dtype=torch.bfloat16,
                                       local_files_only=local_files_only)
        if not getattr(component, 'is_loaded_in_8bit', False):
            raise ValueError('Expected saved bitsandbytes INT8 component: '+name)
        patch_int8_device_moves(component)
        component.to('cpu')
        torch.cuda.empty_cache()
        components[name] = component
    pipe = QwenImage21Pipeline.from_pretrained(model, dtype=torch.bfloat16,
        local_files_only=local_files_only, **components)
    enable_int8_cpu_offload(pipe)
    release_prefix_cache_before_decode(pipe)
    return pipe