diff --git a/modules/sd_hijack_vae.py b/modules/sd_hijack_vae.py index 0b754afb2..cc6b83919 100644 --- a/modules/sd_hijack_vae.py +++ b/modules/sd_hijack_vae.py @@ -27,7 +27,10 @@ def hijack_vae_decode(*args, **kwargs): latents = args[0].to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype if hasattr(shared.sd_model.vae, '_asymmetric_upscale_vae'): res = hijack_vae_upscale(latents, *args[1:], **kwargs) - else: + elif getattr(shared.sd_model, 'sdnext_vae_type', None) == 'Tiny': + from modules.video_models import video_vae + res = video_vae.vae_decode_tiny(latents) # None when the model has no tiny counterpart, and it says so + if res is None: res = shared.sd_model.vae.orig_decode(latents, *args[1:], **kwargs) t1 = time.time() try: diff --git a/modules/vae/sd_vae_taesd.py b/modules/vae/sd_vae_taesd.py index 5ea9eac9c..f40d3b8cc 100644 --- a/modules/vae/sd_vae_taesd.py +++ b/modules/vae/sd_vae_taesd.py @@ -58,8 +58,8 @@ def warn_once(msg, variant=None): def get_model(model_cls, variant=None): - if variant is not None: - pass + if variant is not None: # the caller named the variant it wants; the ladder below only derives one + return model_cls, variant if model_cls in {'sd'}: model_cls = 'sd' variant = shared.opts.taesd_variant diff --git a/modules/video_models/video_save.py b/modules/video_models/video_save.py index 70002f185..028964aa6 100644 --- a/modules/video_models/video_save.py +++ b/modules/video_models/video_save.py @@ -205,7 +205,7 @@ def atomic_save_video( if metadata is None: metadata = {} av = check_av() - if av is None or av is False: + if av is None: log.error('Video: ffmpeg/av not available') return savejob = shared.state.begin('Save video') diff --git a/modules/video_models/video_utils.py b/modules/video_models/video_utils.py index e5dee34dc..1444d4782 100644 --- a/modules/video_models/video_utils.py +++ b/modules/video_models/video_utils.py @@ -54,13 +54,14 @@ def supports_last_frame(model): def check_av(): + """The av module, or None when it is unavailable; callers guard on the None.""" install('av') try: import av av.logging.set_level(av.logging.ERROR) # pylint: disable=c-extension-no-member except Exception as e: log.error(f'av package: {e}') - return False + return None return av diff --git a/modules/video_models/video_vae.py b/modules/video_models/video_vae.py index 5315a6e83..0a5af85e0 100644 --- a/modules/video_models/video_vae.py +++ b/modules/video_models/video_vae.py @@ -4,14 +4,13 @@ from modules.logger import log debug = log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None -vae_type = None +UNIT_RANGE_DECODERS = {'TAEM1'} # taehv shifts its output to [-1,1] on the way out and taem1 leaves it at [0,1] def set_vae_params(p, slicing:bool=True, tiling:bool=True, framewise:bool=True) -> None: - global vae_type # pylint: disable=global-statement - vae_type = p.vae_type if not hasattr(shared.sd_model, 'vae'): return + shared.sd_model.sdnext_vae_type = p.vae_type # the decode hijack reads the choice off the pipe if slicing and hasattr(shared.sd_model.vae, 'enable_slicing'): shared.sd_model.vae.enable_slicing() if (p.frames > p.vae_tile_frames) and (p.vae_tile_frames > 0): @@ -21,33 +20,47 @@ def set_vae_params(p, slicing:bool=True, tiling:bool=True, framewise:bool=True) shared.sd_model.vae.use_framewise_decoding = True if tiling and hasattr(shared.sd_model.vae, 'enable_tiling'): shared.sd_model.vae.enable_tiling() - debug(f'VAE params: type={vae_type} tiling=True frames={p.frames} tile_frames={p.vae_tile_frames} framewise={getattr(shared.sd_model.vae, "use_framewise_decoding", None)}') + debug(f'VAE params: type={p.vae_type} tiling=True frames={p.frames} tile_frames={p.vae_tile_frames} framewise={getattr(shared.sd_model.vae, "use_framewise_decoding", None)}') else: if hasattr(shared.sd_model.vae, 'use_framewise_decoding'): shared.sd_model.vae.use_framewise_decoding = False if hasattr(shared.sd_model.vae, 'disable_tiling'): shared.sd_model.vae.disable_tiling() - debug(f'VAE params: type={vae_type} tiling=False frames={p.frames} tile_frames={p.vae_tile_frames} framewise={getattr(shared.sd_model.vae, "use_framewise_decoding", None)}') + debug(f'VAE params: type={p.vae_type} tiling=False frames={p.frames} tile_frames={p.vae_tile_frames} framewise={getattr(shared.sd_model.vae, "use_framewise_decoding", None)}') def vae_decode_tiny(latents): - if 'Hunyuan' in shared.sd_model.__class__.__name__: + """Decode through the tiny counterpart of the model's vae, or None when there is none to use. + + Returning None leaves the caller on the full vae, so every rejection here is a fallback + rather than a failure. + """ + cls = shared.sd_model.__class__.__name__ + if 'Hunyuan' in cls: variant = 'TAE HunyuanVideo' - elif 'Mochi' in shared.sd_model.__class__.__name__: + elif 'Mochi' in cls: variant = 'TAE MochiVideo' - elif 'WAN' in shared.sd_model.__class__.__name__: + elif 'Wan' in cls: variant = 'TAE WanVideo' - elif 'Kandinsky' in shared.sd_model.__class__.__name__: + elif 'Kandinsky' in cls: variant = 'TAE HunyuanVideo' else: - log.warning(f'Decode: type=Tiny cls={shared.sd_model.__class__.__name__} not supported') + log.warning(f'Decode: type=Tiny cls={cls} not supported') return None from modules.vae import sd_vae_taesd vae, variant = sd_vae_taesd.load_model(variant=variant) if vae is None: return None + expected = getattr(vae, 'latent_channels', None) # 16 on taehv and 12 on taem1, so ask the decoder rather than assume + channels = latents.shape[1] if latents.ndim == 5 else None # the pipes hand the decoder NCTHW + if expected is not None and channels is not None and channels != expected: + log.warning(f'Decode: type=Tiny cls={cls} latents={channels}ch expected={expected}ch not supported') + return None log.debug(f'Decode: type=Tiny cls={vae.__class__.__name__} variant="{variant}" latents={latents.shape}') vae = vae.to(device=devices.device, dtype=devices.dtype) latents = latents.transpose(1, 2).to(device=devices.device, dtype=devices.dtype) - images = vae.decode_video(latents, parallel=False).transpose(1, 2).mul_(2).sub_(1) + images = vae.decode_video(latents, parallel=False).transpose(1, 2) + if type(vae).__name__ in UNIT_RANGE_DECODERS: # the pipelines expect a decode in [-1,1] + images = images.mul_(2).sub_(1) + log.debug(f'Decode: type=Tiny decoded={list(images.shape)} range={images.min().item():.3f}..{images.max().item():.3f}') return (images, None)