"""Build the decoder from the Diffusers H3 VAE config.json alone; decoding is FastH3VAEDecoder's."""
import json
from collections import OrderedDict
from types import SimpleNamespace

from safetensors.torch import load_file

from fast_h3_vae import FastH3VAEDecoder
from noncausal_vit_decoder import NonCausalViTDecoder
from spark_compat import RGBProcessor, make_template
from spark_fa4 import enable_sm120_guard


class SparkFastH3VAEDecoder(FastH3VAEDecoder):
    def __init__(self, vae_config, weights, *, engine='fast', device='cuda', max_engines=1):
        enable_sm120_guard()
        with open(vae_config) as f:
            config = json.load(f)
        model = NonCausalViTDecoder(make_template(config), time_norm=5)
        model.load_state_dict(load_file(str(weights), device='cpu'), strict=True, assign=True)
        assert not any(t.is_meta for t in list(model.parameters()) + list(model.buffers()))
        self.model = model.to(device).eval().requires_grad_(False)
        self.vae = SimpleNamespace(processor=RGBProcessor())
        self.engine, self.max_engines = engine, max_engines
        self._engines = OrderedDict()