| """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() |
|
|
