File size: 1,140 Bytes
775b8fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | """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()
|