"""Diagnostic: what holds 45GB after window 1 encode_video?"""
import os, gc, sys, time
os.environ.setdefault('PYTORCH_ALLOC_CONF', 'expandable_segments:True')
os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'expandable_segments:True')
import torch

sys.path.insert(0, '/workspace/LTX-2/packages/ltx-pipelines/src')
sys.path.insert(0, '/workspace/LTX-2/packages/ltx-core/src')

from ltx_pipelines.ic_lora import ICLoraPipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_core.loader import LoraPathStrengthAndSDOps, LTXV_LORA_COMFY_RENAMING_MAP
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import OffloadMode
import optimum.quanto as oq

CK    = '/workspace/models/LTX-2.3/ltx-2.3-22b-distilled.safetensors'
UP    = '/workspace/models/LTX-2.3/ltx-2.3-spatial-upscaler-x2-1.1.safetensors'
LORA  = '/workspace/models/loras/ltx-2.3-22b-ic-lora-colorization-0.9.safetensors'
GEMMA = '/workspace/models/LTX-2'

print("[DIAG] Pipeline laden...", flush=True)
pipe = ICLoraPipeline(
    distilled_checkpoint_path=CK,
    spatial_upsampler_path=UP,
    gemma_root=GEMMA,
    loras=(LoraPathStrengthAndSDOps(LORA, 1.0, LTXV_LORA_COMFY_RENAMING_MAP),),
    quantization=None,
    compilation_config=None,
    offload_mode=OffloadMode('none'),
)
print(f"[DIAG] Pipeline geladen, VRAM={torch.cuda.memory_allocated()//1024**3}GB", flush=True)

# Prompt cache
with torch.no_grad():
    _ctx = pipe.prompt_encoder(["Cinematic black and white silent film, high contrast"], enhance_first_prompt=False, enhance_prompt_seed=42)[0]
pipe.prompt_encoder = lambda prompts, **kw: [_ctx]
gc.collect(); torch.cuda.empty_cache()
print(f"[DIAG] Gemma freed, VRAM={torch.cuda.memory_allocated()//1024**3}GB", flush=True)

# fp8 DiT
from contextlib import contextmanager
stage = pipe.stage_1
_pre_built_dit = stage._build_transformer()
vram_bf16 = torch.cuda.memory_allocated() // 1024**3
oq.quantize(_pre_built_dit, weights=oq.qfloat8)
# Fix stale refs (TransformerArgsPreprocessor)
_vm = _pre_built_dit.velocity_model
_id_map = {}
for _a in dir(_vm):
    if not _a.startswith('_'):
        _o = getattr(_vm, _a, None)
        if isinstance(_o, torch.nn.Linear):
            _id_map[id(_o)] = _a

def _fix(obj, vm, id_map, depth=0):
    if depth > 6: return
    for k, v in list(getattr(obj, '__dict__', {}).items()):
        if isinstance(v, torch.nn.Linear) and getattr(v, 'weight', True) is None:
            a = id_map.get(id(v))
            if a:
                setattr(obj, k, getattr(vm, a, None))
        elif not isinstance(v, (torch.nn.Module, torch.Tensor)) and hasattr(v, '__dict__'):
            _fix(v, vm, id_map, depth+1)

for _p in [_vm.video_args_preprocessor, getattr(_vm, 'audio_args_preprocessor', None)]:
    if _p: _fix(_p, _vm, _id_map)
oq.freeze(_pre_built_dit)
gc.collect(); torch.cuda.empty_cache()
vram_fp8 = torch.cuda.memory_allocated() // 1024**3
print(f"[DIAG] DiT fp8: {vram_fp8}GB (was bf16 {vram_bf16}GB)", flush=True)

@contextmanager
def _ctx_mgr(**kw):
    yield _pre_built_dit
    gc.collect(); torch.cuda.empty_cache()

stage._transformer_ctx = _ctx_mgr
pipe.stage_2._transformer_ctx = _ctx_mgr

# Tiled encode patch
import ltx_pipelines.iclora_utils as _icu
import ltx_pipelines.ic_lora as _icl
from ltx_core.model.video_vae.tiling import SpatialTilingConfig as _STC, TemporalTilingConfig as _TTC
_TILE = TilingConfig(spatial_config=_STC(512, 64), temporal_config=_TTC(80, 24))
_orig = _icu.append_ic_lora_reference_video_conditionings
_icu.append_ic_lora_reference_video_conditionings = lambda *a, tiling_config=None, **kw: _orig(*a, tiling_config=_TILE, **kw)
_icl.append_ic_lora_reference_video_conditionings = _icu.append_ic_lora_reference_video_conditionings

# Run window 0
import json
with open('/workspace/ltx_co_manifest.json') as f:
    w = json.load(f)[0]

tiling = TilingConfig.default()
imgs = [ImageConditioningInput(path=w['keyframe'], frame_idx=w.get('kf_pos', 40), strength=1.0)]
print(f"[DIAG] Render venster 0...", flush=True)
with torch.no_grad():
    video, audio = pipe(
        prompt=w['prompt'], seed=42, height=1088, width=1920, num_frames=81, frame_rate=24,
        images=imgs, video_conditioning=[(w['gray'], 1.0)], tiling_config=tiling,
        conditioning_attention_strength=1.0, skip_stage_2=True, conditioning_attention_mask=None,
    )
print(f"[DIAG] Pipe klaar, VRAM={torch.cuda.memory_allocated()//1024**3}GB", flush=True)

import os
os.makedirs(os.path.dirname(w['out']), exist_ok=True)
encode_video(video=video, fps=24, audio=audio, output_path=w['out'],
             video_chunks_number=get_video_chunks_number(81, tiling))
print(f"[DIAG] encode_video klaar, VRAM={torch.cuda.memory_allocated()//1024**3}GB", flush=True)

del video, audio
gc.collect(); torch.cuda.empty_cache()
print(f"[DIAG] Na del+gc, VRAM={torch.cuda.memory_allocated()//1024**3}GB", flush=True)

# DIAGNOSTIEK: wat houdt CUDA memory?
print(f"\n[DIAG] === CUDA GEHEUGEN ANALYSE ===", flush=True)
print(f"allocated={torch.cuda.memory_allocated()//1024**3}GB reserved={torch.cuda.memory_reserved()//1024**3}GB", flush=True)

# Scan pipe.video_decoder
print("\n[DIAG] pipe.video_decoder module scan:", flush=True)
vd = pipe.video_decoder
total_vd = 0
for name, m in vd.named_modules():
    # Parameters
    p_sizes = [(n, p.numel()*p.element_size()//1024**2) for n,p in m.named_parameters(recurse=False) if p is not None and p.is_cuda]
    b_sizes = [(n, b.numel()*b.element_size()//1024**2) for n,b in m.named_buffers(recurse=False) if b is not None and b.is_cuda]
    # Also scan __dict__ for cached tensors
    dict_sizes = []
    for k, v in m.__dict__.items():
        if isinstance(v, torch.Tensor) and v.is_cuda:
            sz = v.numel() * v.element_size() // 1024**2
            dict_sizes.append((k, sz))
        elif isinstance(v, list):
            for item in v:
                if isinstance(item, torch.Tensor) and item.is_cuda:
                    sz = item.numel() * item.element_size() // 1024**2
                    dict_sizes.append((f'{k}[]', sz))
    total_m = sum(s for _,s in p_sizes+b_sizes+dict_sizes)
    if total_m > 100:  # > 100MB
        total_vd += total_m
        print(f"  {name or 'root'}: {total_m}MB", flush=True)
        for n, s in p_sizes+b_sizes+dict_sizes:
            if s > 10:
                print(f"    .{n}: {s}MB", flush=True)

print(f"\n[DIAG] Total in video_decoder: ~{total_vd}MB", flush=True)

# Scan pipe.stage_1 (NOT the DiT)
print("\n[DIAG] pipe.stage_1 non-DiT scan:", flush=True)
stage1 = pipe.stage_1
for attr in dir(stage1):
    if attr.startswith('_'): continue
    try:
        v = getattr(stage1, attr)
        if isinstance(v, torch.Tensor) and v.is_cuda:
            print(f"  stage_1.{attr}: {v.numel()*v.element_size()//1024**2}MB shape={list(v.shape)}", flush=True)
        elif isinstance(v, torch.nn.Module):
            sz = sum(p.numel()*p.element_size() for p in v.parameters() if p.is_cuda) // 1024**2
            if sz > 100:
                print(f"  stage_1.{attr}: ~{sz}MB (nn.Module)", flush=True)
    except Exception:
        pass

print("\n[DIAG] KLAAR", flush=True)