From 5c46904fd869a7aae9384c031d342b2b78db4758 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 20 Mar 2025 14:31:30 -0400 Subject: [PATCH] video tab alpha releasee Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 50 +++-- modules/lora/network_overrides.py | 8 +- modules/modeldata.py | 15 +- modules/processing_vae.py | 2 + modules/sd_offload.py | 6 +- modules/sd_vae_remote.py | 18 +- modules/sd_vae_taesd.py | 55 ++++-- modules/taesd/taehv.py | 284 ++++++++++++++++++++++++++++ modules/taesd/taem1.py | 272 ++++++++++++++++++++++++++ modules/ui_video.py | 6 +- modules/video_models/models_def.py | 33 ++++ modules/video_models/run_allegro.py | 8 +- modules/video_models/run_cog.py | 8 +- modules/video_models/run_hunyuan.py | 13 +- modules/video_models/run_ltx.py | 8 +- modules/video_models/run_mochi.py | 8 +- modules/video_models/video_load.py | 80 ++++++++ modules/video_models/video_utils.py | 102 +--------- modules/video_models/video_vae.py | 62 ++++++ scripts/cogvideo.py | 8 +- 20 files changed, 874 insertions(+), 172 deletions(-) create mode 100644 modules/taesd/taehv.py create mode 100644 modules/taesd/taem1.py create mode 100644 modules/video_models/video_load.py create mode 100644 modules/video_models/video_vae.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e3cac31c..f784e5f56 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,21 +2,49 @@ ## Update for 2025-03-19 -### TODO - - Gemma3 requires `transformers==git+https://github.com/huggingface/transformers@v4.49.0-Gemma-3` - - Remote VAE encode for SD15 and Flux.1: - - HunyuanVideo-I2V: - - HunyuanVideo: Remote VAE - - HunyuanVideo: Tiny VAE - - LTXVideo-095: Condition input - - LTXVideo-095: Broken offloading +### ToDo/Limitations -### Highlights for 2025-03-17 + - VLM Gemma3: requires `transformers==git+https://github.com/huggingface/transformers@v4.49.0-Gemma-3` + - VAE Remote encode: SD15 and Flux.1 issues: + - Video: ModernUI support is TBD + - Video: API support is TBD + - Video: Wiki page is TBD + - Video: HunyuanVideo-I2V incompatible with latest transformers + - Video: LTXVideo-095 support for conditioned input + - Video: LTXVideo-095 support for offloading -Support for [CogView 4](https://huggingface.co/THUDM/CogView4-6B), new CLiP models, improvements to remote VAE, additional docs/guides. +### Highlights for 2025-03-20 -### Details for 2025-03-17 +Brand new Video processing module with support for all latest models: **WAN21, Hunyuan, LTX, Cog, Allegro, Mochi1** and more! +Plus support for CogView-4, new CLiP models, improvements to remote VAE, additional docs/guides. +### Details for 2025-03-20 + +- **Video tab** + - initial release so consider this as alpha version + - new top-level tab, replaces previous *video* script in text/image tabs + old scripts are still present, but will be removed in the future + - support for all latest models: + - [Hunyuan](https://huggingface.co/Tencent/HunyuanVideo): *HunyuanVideo, FastHunyuan, SkyReels* | *T2V, I2V* + - [WAN21](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers): *1.3B, 13B* | *T2V, I2V* + - [LTXVideo](https://huggingface.co/Lightricks/LTX-Video): *0.9.0, 0.9.1, 0.9.5* | *T2V, I2V* + - [CogVideoX](https://huggingface.co/THUDM/CogVideoX-5b): *2B, 5B* | *T2V, I2V* + - [Allegro](https://huggingface.co/rhymes-ai/Allegro): *T2V* + - [Mochi1](https://huggingface.co/genmo/mochi-1-preview): *T2V* + - decoding: + - **Default**: use vae from model + - **Tiny VAE**: support for *Hunyuan, WAN, Mochi* + - **Remote VAE**: support for *Hunyuan* + - **LoRA**: support for *Hunyuan, LTX, WAN, Mochi, Cog* + - additional key points: + - all models are auto-downloaded upon first use + - optional video interpolation while creating video files + - optional video preview in ui + - support for balanced offloading and model offloading + - on-the-fly quantization: *BnB, Quanto, TorchAO* + - different video models support different video resolutions, frame counts, etc. + and may require specific settings - see model links for details + - see *ToDo/Limitations* section for additional notes - **Models** - [THUDM CogView 4 6B](https://huggingface.co/THUDM/CogView4-6B) new foundation model for image generation based o GLM-4 text encoder and a flow-based diffusion transformer diff --git a/modules/lora/network_overrides.py b/modules/lora/network_overrides.py index 65448ef2d..22d251c47 100644 --- a/modules/lora/network_overrides.py +++ b/modules/lora/network_overrides.py @@ -30,8 +30,14 @@ force_models = [ # forced always 'sc', 'kandinsky', 'hunyuandit', - 'hunyuanvideo', 'auraflow', + # video models + 'hunyuanvideo', + 'cogvideo', + 'wanvideo', + 'ltxvideo', + 'mochivideo', + 'allegrovideo', ] force_classes = [ # forced always diff --git a/modules/modeldata.py b/modules/modeldata.py index 078a5b372..f066eaa68 100644 --- a/modules/modeldata.py +++ b/modules/modeldata.py @@ -29,8 +29,6 @@ def get_model_type(pipe): model_type = 'auraflow' elif "Flux" in name: model_type = 'f1' - elif "Mochi" in name: - model_type = 'mochi' elif "Lumina2" in name: model_type = 'lumina2' elif "Lumina" in name: @@ -41,12 +39,21 @@ def get_model_type(pipe): model_type = 'cogview3' elif "CogView4" in name: model_type = 'cogview4' - elif "CogVideo" in name: - model_type = 'cogvideox' elif "Sana" in name: model_type = 'sana' + # video models + elif "CogVideo" in name: + model_type = 'cogvideo' elif 'HunyuanVideoPipeline' in name or 'HunyuanSkyreels' in name: model_type = 'hunyuanvideo' + elif 'Wan' in name: + model_type = 'wanvideo' + elif 'LTX' in name: + model_type = 'ltxvideo' + elif "Mochi" in name: + model_type = 'mochivideo' + elif "Allegro" in name: + model_type = 'allegrovideo' else: model_type = name return model_type diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 348bbc0c5..00eecc092 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -239,6 +239,8 @@ def vae_postprocess(tensor, model, output_type='np'): if len(tensor.shape) == 3 and tensor.shape[0] == 3: tensor = tensor.unsqueeze(0) if hasattr(model, 'video_processor'): + if len(tensor.shape) == 6 and tensor.shape[1] == 1: + tensor = tensor.squeeze(0) images = model.video_processor.postprocess_video(tensor, output_type='pil') elif hasattr(model, 'image_processor'): images = model.image_processor.postprocess(tensor, output_type=output_type) diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 2c4126209..cadfd3c02 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -10,7 +10,7 @@ from modules.timer import process as process_timer debug_move = shared.log.trace if os.environ.get('SD_MOVE_DEBUG', None) is not None else lambda *args, **kwargs: None -should_offload = ['sc', 'sd3', 'f1', 'hunyuandit', 'auraflow', 'omnigen', 'hunyuanvideo', 'cogvideox', 'mochi', 'cogview4'] +should_offload = ['sc', 'sd3', 'f1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4'] offload_hook_instance = None balanced_offload_exclude = ['OmniGenPipeline', 'CogView4Pipeline'] @@ -66,7 +66,7 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False): if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate): sd_model.has_accelerate = False if shared.opts.diffusers_offload_mode == "none": - if shared.sd_model_type in should_offload: + if shared.sd_model_type in should_offload or 'video' in shared.sd_model_type: shared.log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model') else: shared.log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') @@ -183,7 +183,7 @@ class OffloadHook(accelerate.hooks.ModelHook): return module -def apply_balanced_offload(sd_model, exclude=[]): +def apply_balanced_offload(sd_model=None, exclude=[]): global offload_hook_instance # pylint: disable=global-statement if shared.opts.diffusers_offload_mode != "balanced": return sd_model diff --git a/modules/sd_vae_remote.py b/modules/sd_vae_remote.py index c3591af7b..3c1846ea8 100644 --- a/modules/sd_vae_remote.py +++ b/modules/sd_vae_remote.py @@ -38,11 +38,14 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_ return tensors t0 = time.time() modelloader.hf_login() - latents = latents.unsqueeze(0) if len(latents.shape) == 3 else latents + latent_copy = latents.detach().clone().to(device=devices.cpu, dtype=devices.dtype) + latent_copy = latents.unsqueeze(0) if len(latents.shape) == 3 else latents + if model_type == 'hunyuanvideo': + latent_copy = latent_copy.unsqueeze(0) - for i in range(latents.shape[0]): + for i in range(latent_copy.shape[0]): try: - latent = latents[i].detach().clone().to(device=devices.cpu, dtype=devices.dtype) + latent = latent_copy[i] if model_type != 'f1': latent = latent.unsqueeze(0) params = { @@ -51,7 +54,12 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_ "dtype": str(latent.dtype).split(".", maxsplit=1)[-1], } headers = { "Content-Type": "tensor/binary" } - if shared.opts.remote_vae_type == 'png': + if 'video' in model_type: + params["partial_postprocess"] = False + params["output_type"] = "pt" + params["output_tensor_type"] = "binary" + headers["Accept"] = "tensor/binary" + elif shared.opts.remote_vae_type == 'png': params["image_format"] = "png" params["output_type"] = "pil" headers["Accept"] = "image/png" @@ -81,7 +89,7 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_ shared.log.error(f'Decode: type="remote" model={model_type} code={response.status_code} shape={latent.shape} url="{url}" args={params} headers={response.headers} response={response.json()}') else: content += len(response.content) - if shared.opts.remote_vae_type == 'raw': + if shared.opts.remote_vae_type == 'raw' or 'video' in model_type: shape = json.loads(response.headers["shape"]) dtype = response.headers["dtype"] tensor = torch.frombuffer(bytearray(response.content), dtype=dtypes[dtype]).reshape(shape) diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index c8a1b882f..a2a447a3e 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -16,6 +16,9 @@ TAESD_MODELS = { 'TAESD 1.2 Chocolate-Dipped Shortbread': { 'fn': 'taesd_12_', 'uri': 'https://github.com/madebyollin/taesd/raw/8909b44e3befaa0efa79c5791e4fe1c4d4f7884e', 'model': None }, 'TAESD 1.1 Fruit Loops': { 'fn': 'taesd_11_', 'uri': 'https://github.com/madebyollin/taesd/raw/3e8a8a2ab4ad4079db60c1c7dc1379b4cc0c6b31', 'model': None }, 'TAESD 1.0': { 'fn': 'taesd_10_', 'uri': 'https://github.com/madebyollin/taesd/raw/88012e67cf0454e6d90f98911fe9d4aef62add86', 'model': None }, + 'TAE HunyuanVideo': { 'fn': 'taehv.pth', 'uri': 'https://github.com/madebyollin/taehv/raw/refs/heads/main/taehv.pth', 'model': None }, + 'TAE WanVideo': { 'fn': 'taew1.pth', 'uri': 'https://github.com/madebyollin/taehv/raw/refs/heads/main/taew2_1.pth', 'model': None }, + 'TAE MochiVideo': { 'fn': 'taem1.pth', 'uri': 'https://github.com/madebyollin/taem1/raw/refs/heads/main/taem1.pth', 'model': None }, } CQYAN_MODELS = { 'Hybrid-Tiny SD': { @@ -35,49 +38,63 @@ prev_model = '' lock = threading.Lock() -def warn_once(msg): +def warn_once(msg, variant=None): from modules import shared + variant = variant or shared.opts.taesd_variant global prev_warnings # pylint: disable=global-statement if not prev_warnings: prev_warnings = True - shared.log.error(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}": {msg}') + shared.log.error(f'Decode: type="taesd" variant="{variant}": {msg}') return Image.new('RGB', (8, 8), color = (0, 0, 0)) -def get_model(model_type = 'decoder'): +def get_model(model_type = 'decoder', variant = None): global prev_cls, prev_type, prev_model # pylint: disable=global-statement from modules import shared cls = shared.sd_model_type if cls == 'ldm': cls = 'sd' + variant = variant or shared.opts.taesd_variant folder = os.path.join(paths.models_path, "TAESD") os.makedirs(folder, exist_ok=True) - if 'sd' not in cls and 'f1' not in cls: + if 'video' in cls: + return None + if ('sd' not in cls) and ('f1' not in cls): warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported') return None - if shared.opts.taesd_variant.startswith('TAESD'): - cfg = TAESD_MODELS[shared.opts.taesd_variant] - if (cls == prev_cls) and (model_type == prev_type) and (shared.opts.taesd_variant == prev_model) and (cfg['model'] is not None): + if variant.startswith('TAESD'): + cfg = TAESD_MODELS[variant] + if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None): return cfg['model'] fn = os.path.join(folder, cfg['fn'] + cls + '_' + model_type + '.pth') if not os.path.exists(fn): uri = cfg['uri'] + '/tae' + cls + '_' + model_type + '.pth' try: - shared.log.info(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}": uri="{uri}" fn="{fn}" download') + shared.log.info(f'Decode: type="taesd" variant="{variant}": uri="{uri}" fn="{fn}" download') torch.hub.download_url_to_file(uri, fn) except Exception as e: warn_once(f'download uri={uri} {e}') if os.path.exists(fn): prev_cls = cls prev_type = model_type - prev_model = shared.opts.taesd_variant - shared.log.debug(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}" fn="{fn}" load') - from modules.taesd.taesd import TAESD - TAESD_MODELS[shared.opts.taesd_variant]['model'] = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None) - return TAESD_MODELS[shared.opts.taesd_variant]['model'] - elif shared.opts.taesd_variant.startswith('Hybrid'): - cfg = CQYAN_MODELS[shared.opts.taesd_variant].get(cls, None) - if (cls == prev_cls) and (model_type == prev_type) and (shared.opts.taesd_variant == prev_model) and (cfg['model'] is not None): + prev_model = variant + shared.log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" load') + if 'TAEHV' in variant: + from modules.taesd.taehv import TAEHV + TAESD_MODELS[variant]['model'] = TAEHV(checkpoint_path=fn) + if 'TAEW2' in variant: + from modules.taesd.taehv import TAEHV + TAESD_MODELS[variant]['model'] = TAEHV(checkpoint_path=fn) + elif 'TAEM1' in variant: + from modules.taesd.taem1 import TAEM1 + TAESD_MODELS[variant]['model'] = TAEM1(checkpoint_path=fn) + else: + from modules.taesd.taesd import TAESD + TAESD_MODELS[variant]['model'] = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None) + return TAESD_MODELS[variant]['model'] + elif variant.startswith('Hybrid'): + cfg = CQYAN_MODELS[variant].get(cls, None) + if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None): return cfg['model'] if cfg is None: warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported') @@ -85,8 +102,8 @@ def get_model(model_type = 'decoder'): repo = cfg['repo'] prev_cls = cls prev_type = model_type - prev_model = shared.opts.taesd_variant - shared.log.debug(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}" id="{repo}" load') + prev_model = variant + shared.log.debug(f'Decode: type="taesd" variant="{variant}" id="{repo}" load') dtype = devices.dtype_vae if devices.dtype_vae != torch.bfloat16 else torch.float16 # taesd does not support bf16 if 'tiny' in repo: from diffusers.models import AutoencoderTiny @@ -95,7 +112,7 @@ def get_model(model_type = 'decoder'): from modules.taesd.hybrid_small import AutoencoderSmall vae = AutoencoderSmall.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir, torch_dtype=dtype) vae = vae.to(devices.device, dtype=dtype) - CQYAN_MODELS[shared.opts.taesd_variant][cls]['model'] = vae + CQYAN_MODELS[variant][cls]['model'] = vae return vae else: warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported') diff --git a/modules/taesd/taehv.py b/modules/taesd/taehv.py new file mode 100644 index 000000000..4a424f137 --- /dev/null +++ b/modules/taesd/taehv.py @@ -0,0 +1,284 @@ +#!/usr/bin/env python3 +""" +Tiny AutoEncoder for Hunyuan Video +(DNN for encoding / decoding videos to Hunyuan Video's latent space) +""" +from collections import namedtuple +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm.auto import tqdm + +DecoderResult = namedtuple("DecoderResult", ("frame", "memory")) +TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index")) + +def conv(n_in, n_out, **kwargs): + return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) + +class Clamp(nn.Module): + def forward(self, x): + return torch.tanh(x / 3) * 3 + +class MemBlock(nn.Module): + def __init__(self, n_in, n_out): + super().__init__() + self.conv = nn.Sequential(conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out)) + self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() + self.act = nn.ReLU(inplace=True) + def forward(self, x, past): + return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) + +class TPool(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f*stride,n_f, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + return self.conv(x.reshape(-1, self.stride * C, H, W)) + +class TGrow(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f, n_f*stride, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + x = self.conv(x) + return x.reshape(-1, C, H, W) + +def apply_model_with_memblocks(model, x, parallel, show_progress_bar): + """ + Apply a sequential model with memblocks to the given input. + Args: + - model: nn.Sequential of blocks to apply + - x: input data, of dimensions NTCHW + - parallel: if True, parallelize over timesteps (fast but uses O(T) memory) + if False, each timestep will be processed sequentially (slow but uses O(1) memory) + - show_progress_bar: if True, enables tqdm progressbar display + + Returns NTCHW tensor of output data. + """ + assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor" + N, T, C, H, W = x.shape + if parallel: + x = x.reshape(N*T, C, H, W) + # parallel over input timesteps, iterate over blocks + for b in tqdm(model, disable=not show_progress_bar): + if isinstance(b, MemBlock): + NT, C, H, W = x.shape + T = NT // N + _x = x.reshape(N, T, C, H, W) + mem = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape) + x = b(x, mem) + else: + x = b(x) + NT, C, H, W = x.shape + T = NT // N + x = x.view(N, T, C, H, W) + else: + # TODO(oboerbohan): at least on macos this still gradually uses more memory during decode... + # need to fix :( + out = [] + # iterate over input timesteps and also iterate over blocks. + # because of the cursed TPool/TGrow blocks, this is not a nested loop, + # it's actually a ***graph traversal*** problem! so let's make a queue + work_queue = [TWorkItem(xt, 0) for t, xt in enumerate(x.reshape(N, T * C, H, W).chunk(T, dim=1))] + # in addition to manually managing our queue, we also need to manually manage our progressbar. + # we'll update it for every source node that we consume. + progress_bar = tqdm(range(T), disable=not show_progress_bar) + # we'll also need a separate addressable memory per node as well + mem = [None] * len(model) + while work_queue: + xt, i = work_queue.pop(0) + if i == 0: + # new source node consumed + progress_bar.update(1) + if i == len(model): + # reached end of the graph, append result to output list + out.append(xt) + else: + # fetch the block to process + b = model[i] + if isinstance(b, MemBlock): + # mem blocks are simple since we're visiting the graph in causal order + if mem[i] is None: + xt_new = b(xt, xt * 0) + mem[i] = xt + else: + xt_new = b(xt, mem[i]) + mem[i].copy_(xt) # inplace might reduce mysterious pytorch memory allocations? doesn't help though + # add successor to work queue + work_queue.insert(0, TWorkItem(xt_new, i+1)) + elif isinstance(b, TPool): + # pool blocks are miserable + if mem[i] is None: + mem[i] = [] # pool memory is itself a queue of inputs to pool + mem[i].append(xt) + if len(mem[i]) > b.stride: + # pool mem is in invalid state, we should have pooled before this + raise ValueError("???") + elif len(mem[i]) < b.stride: + # pool mem is not yet full, go back to processing the work queue + pass + else: + # pool mem is ready, run the pool block + N, C, H, W = xt.shape + xt = b(torch.cat(mem[i], 1).view(N*b.stride, C, H, W)) + # reset the pool mem + mem[i] = [] + # add successor to work queue + work_queue.insert(0, TWorkItem(xt, i+1)) + elif isinstance(b, TGrow): + xt = b(xt) + NT, C, H, W = xt.shape + # each tgrow has multiple successor nodes + for xt_next in reversed(xt.view(N, b.stride*C, H, W).chunk(b.stride, 1)): + # add successor to work queue + work_queue.insert(0, TWorkItem(xt_next, i+1)) + else: + # normal block with no funny business + xt = b(xt) + # add successor to work queue + work_queue.insert(0, TWorkItem(xt, i+1)) + progress_bar.close() + x = torch.stack(out, 1) + return x + +class TAEHV(nn.Module): + latent_channels = 16 + image_channels = 3 + def __init__(self, checkpoint_path="taehv.pth", decoder_time_upscale=(True, True), decoder_space_upscale=(True, True, True)): + """Initialize pretrained TAEHV from the given checkpoint. + + Arg: + checkpoint_path: path to weight file to load. taehv.pth for Hunyuan, taew2_1.pth for Wan 2.1. + decoder_time_upscale: whether temporal upsampling is enabled for each block. upsampling can be disabled for a cheaper preview. + decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview. + """ + super().__init__() + self.encoder = nn.Sequential( + conv(TAEHV.image_channels, 64), nn.ReLU(inplace=True), + TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + TPool(64, 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + conv(64, TAEHV.latent_channels), + ) + n_f = [256, 128, 64, 64] + self.frames_to_trim = 2**sum(decoder_time_upscale) - 1 + self.decoder = nn.Sequential( + Clamp(), conv(TAEHV.latent_channels, n_f[0]), nn.ReLU(inplace=True), + MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False), + MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False), + MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), conv(n_f[2], n_f[3], bias=False), + nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels), + ) + if checkpoint_path is not None: + self.load_state_dict(self.patch_tgrow_layers(torch.load(checkpoint_path, map_location="cpu", weights_only=True))) + + def patch_tgrow_layers(self, sd): + """Patch TGrow layers to use a smaller kernel if needed. + + Args: + sd: state dict to patch + """ + new_sd = self.state_dict() + for i, layer in enumerate(self.decoder): + if isinstance(layer, TGrow): + key = f"decoder.{i}.conv.weight" + if sd[key].shape[0] > new_sd[key].shape[0]: + # take the last-timestep output channels + sd[key] = sd[key][-new_sd[key].shape[0]:] + return sd + + def encode_video(self, x, parallel=True, show_progress_bar=True): + """Encode a sequence of frames. + + Args: + x: input NTCHW RGB (C=3) tensor with values in [0, 1]. + parallel: if True, all frames will be processed at once. + (this is faster but may require more memory). + if False, frames will be processed sequentially. + Returns NTCHW latent tensor with ~Gaussian values. + """ + return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar) + + def decode_video(self, x, parallel=True, show_progress_bar=True): + """Decode a sequence of frames. + + Args: + x: input NTCHW latent (C=12) tensor with ~Gaussian values. + parallel: if True, all frames will be processed at once. + (this is faster but may require more memory). + if False, frames will be processed sequentially. + Returns NTCHW RGB tensor with ~[0, 1] values. + """ + x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar) + return x[:, self.frames_to_trim:] + + def forward(self, x): + return self.c(x) + +@torch.no_grad() +def main(): + """Run TAEHV roundtrip reconstruction on the given video paths.""" + import sys + import cv2 # no highly esteemed deed is commemorated here + + class VideoTensorReader: + def __init__(self, video_file_path): + self.cap = cv2.VideoCapture(video_file_path) + assert self.cap.isOpened(), f"Could not load {video_file_path}" + self.fps = self.cap.get(cv2.CAP_PROP_FPS) + def __iter__(self): + return self + def __next__(self): + ret, frame = self.cap.read() + if not ret: + self.cap.release() + raise StopIteration # End of video or error + return torch.from_numpy(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).permute(2, 0, 1) # BGR HWC -> RGB CHW + + class VideoTensorWriter: + def __init__(self, video_file_path, width_height, fps=30): + self.writer = cv2.VideoWriter(video_file_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, width_height) + assert self.writer.isOpened(), f"Could not create writer for {video_file_path}" + def write(self, frame_tensor): + assert frame_tensor.ndim == 3 and frame_tensor.shape[0] == 3, f"{frame_tensor.shape}??" + self.writer.write(cv2.cvtColor(frame_tensor.permute(1, 2, 0).numpy(), cv2.COLOR_RGB2BGR)) # RGB CHW -> BGR HWC + def __del__(self): + if hasattr(self, 'writer'): + self.writer.release() + + dev = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu") + dtype = torch.float16 + print("Using device", dev, "and dtype", dtype) + taehv = TAEHV().to(dev, dtype) + for video_path in sys.argv[1:]: + print(f"Processing {video_path}...") + video_in = VideoTensorReader(video_path) + video = torch.stack(list(video_in), 0)[None] + vid_dev = video.to(dev, dtype).div_(255.0) + # convert to device tensor + if video.numel() < 100_000_000: + print(f" {video_path} seems small enough, will process all frames in parallel") + # convert to device tensor + vid_enc = taehv.encode_video(vid_dev) + print(f" Encoded {video_path} -> {vid_enc.shape}. Decoding...") + vid_dec = taehv.decode_video(vid_enc) + print(f" Decoded {video_path} -> {vid_dec.shape}") + else: + print(f" {video_path} seems large, will process each frame sequentially") + # convert to device tensor + vid_enc = taehv.encode_video(vid_dev, parallel=False) + print(f" Encoded {video_path} -> {vid_enc.shape}. Decoding...") + vid_dec = taehv.decode_video(vid_enc, parallel=False) + print(f" Decoded {video_path} -> {vid_dec.shape}") + video_out_path = video_path + ".reconstructed_by_taehv.mp4" + video_out = VideoTensorWriter(video_out_path, (vid_dec.shape[-1], vid_dec.shape[-2]), fps=int(round(video_in.fps))) + for frame in vid_dec.clamp_(0, 1).mul_(255).round_().byte().cpu()[0]: + video_out.write(frame) + print(f" Saved to {video_out_path}") + +if __name__ == "__main__": + main() diff --git a/modules/taesd/taem1.py b/modules/taesd/taem1.py new file mode 100644 index 000000000..7d59ca2b6 --- /dev/null +++ b/modules/taesd/taem1.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 +""" +Tiny AutoEncoder for Mochi 1 +(DNN for encoding / decoding videos to Mochi 1's latent space) +""" +from collections import namedtuple +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm.auto import tqdm + +DecoderResult = namedtuple("DecoderResult", ("frame", "memory")) +TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index")) + +def conv(n_in, n_out, **kwargs): + return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) + +class Clamp(nn.Module): + def forward(self, x): + return torch.tanh(x / 3) * 3 + +class MemBlock(nn.Module): + def __init__(self, n_in, n_out): + super().__init__() + self.conv = nn.Sequential(conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out)) + self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() + self.act = nn.ReLU(inplace=True) + def forward(self, x, past): + return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) + +class TPool(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f*stride,n_f, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + return self.conv(x.reshape(-1, self.stride * C, H, W)) + +class TGrow(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f, n_f*stride, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + x = self.conv(x) + return x.reshape(-1, C, H, W) + +def apply_model_with_memblocks(model, x, parallel, show_progress_bar): + """ + Apply a sequential model with memblocks to the given input. + Args: + - model: nn.Sequential of blocks to apply + - x: input data, of dimensions NTCHW + - parallel: if True, parallelize over timesteps (fast but uses O(T) memory) + if False, each timestep will be processed sequentially (slow but uses O(1) memory) + - show_progress_bar: if True, enables tqdm progressbar display + + Returns NTCHW tensor of output data. + """ + assert x.ndim == 5, f"TAEM1 operates on NTCHW tensors, but got {x.ndim}-dim tensor" + N, T, C, H, W = x.shape + if parallel: + x = x.reshape(N*T, C, H, W) + # parallel over input timesteps, iterate over blocks + for b in tqdm(model, disable=not show_progress_bar): + if isinstance(b, MemBlock): + NT, C, H, W = x.shape + T = NT // N + _x = x.reshape(N, T, C, H, W) + mem = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape) + x = b(x, mem) + else: + x = b(x) + NT, C, H, W = x.shape + T = NT // N + x = x.view(N, T, C, H, W) + else: + # TODO(oboerbohan): at least on macos this still gradually uses more memory during decode... + # need to fix :( + out = [] + # iterate over input timesteps and also iterate over blocks. + # because of the cursed TPool/TGrow blocks, this is not a nested loop, + # it's actually a ***graph traversal*** problem! so let's make a queue + work_queue = [TWorkItem(xt, 0) for t, xt in enumerate(x.reshape(N, T * C, H, W).chunk(T, dim=1))] + # in addition to manually managing our queue, we also need to manually manage our progressbar. + # we'll update it for every source node that we consume. + progress_bar = tqdm(range(T), disable=not show_progress_bar) + # we'll also need a separate addressable memory per node as well + mem = [None] * len(model) + while work_queue: + xt, i = work_queue.pop(0) + if i == 0: + # new source node consumed + progress_bar.update(1) + if i == len(model): + # reached end of the graph, append result to output list + out.append(xt) + else: + # fetch the block to process + b = model[i] + if isinstance(b, MemBlock): + # mem blocks are simple since we're visiting the graph in causal order + if mem[i] is None: + xt_new = b(xt, xt * 0) + mem[i] = xt + else: + xt_new = b(xt, mem[i]) + mem[i].copy_(xt) # inplace might reduce mysterious pytorch memory allocations? doesn't help though + # add successor to work queue + work_queue.insert(0, TWorkItem(xt_new, i+1)) + elif isinstance(b, TPool): + # pool blocks are miserable + if mem[i] is None: + mem[i] = [] # pool memory is itself a queue of inputs to pool + mem[i].append(xt) + if len(mem[i]) > b.stride: + # pool mem is in invalid state, we should have pooled before this + raise ValueError("???") + elif len(mem[i]) < b.stride: + # pool mem is not yet full, go back to processing the work queue + pass + else: + # pool mem is ready, run the pool block + N, C, H, W = xt.shape + xt = b(torch.cat(mem[i], 1).view(N*b.stride, C, H, W)) + # reset the pool mem + mem[i] = [] + # add successor to work queue + work_queue.insert(0, TWorkItem(xt, i+1)) + elif isinstance(b, TGrow): + xt = b(xt) + NT, C, H, W = xt.shape + # each tgrow has multiple successor nodes + for xt_next in reversed(xt.view(N, b.stride*C, H, W).chunk(b.stride, 1)): + # add successor to work queue + work_queue.insert(0, TWorkItem(xt_next, i+1)) + else: + # normal block with no funny business + xt = b(xt) + # add successor to work queue + work_queue.insert(0, TWorkItem(xt, i+1)) + progress_bar.close() + x = torch.stack(out, 1) + return x + +class TAEM1(nn.Module): + latent_channels = 12 + image_channels = 3 + def __init__(self, checkpoint_path="taem1.pth"): + """Initialize pretrained TAEM1 from the given checkpoints.""" + super().__init__() + self.encoder = nn.Sequential( + conv(TAEM1.image_channels, 64), nn.ReLU(inplace=True), + TPool(64, 3), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + TPool(64, 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + conv(64, TAEM1.latent_channels), + ) + n_f = [256, 128, 64, 64] + self.decoder = nn.Sequential( + Clamp(), conv(TAEM1.latent_channels, n_f[0]), nn.ReLU(inplace=True), + MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False), + MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2), TGrow(n_f[1], 2), conv(n_f[1], n_f[2], bias=False), + MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2), TGrow(n_f[2], 3), conv(n_f[2], n_f[3], bias=False), + nn.ReLU(inplace=True), conv(n_f[3], TAEM1.image_channels), + ) + if checkpoint_path is not None: + self.load_state_dict(torch.load(checkpoint_path, map_location="cpu", weights_only=True)) + + def encode_video(self, x, parallel=True, show_progress_bar=True): + """Encode a sequence of frames. + + Args: + x: input NTCHW RGB (C=3) tensor with values in [0, 1]. + parallel: if True, all frames will be processed at once. + (this is faster but may require more memory). + if False, frames will be processed sequentially. + Returns NTCHW latent tensor with ~Gaussian values. + """ + return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar) + + def decode_video(self, x, parallel=True, show_progress_bar=True): + """Decode a sequence of frames. + + Args: + x: input NTCHW latent (C=12) tensor with ~Gaussian values. + parallel: if True, all frames will be processed at once. + (this is faster but may require more memory). + if False, frames will be processed sequentially. + Returns NTCHW RGB tensor with ~[0, 1] values. + """ + x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar) + # NOTE: + # the Mochi VAE does not preserve shape along the time axis; + # videos are encoded to floor((n_in - 1)/6)+1 latent frames + # (which makes sense, it's stride 6, so 12 -> 2 and 13->3) + # but then they're decoded to only the *minimal* number + # of input frames (3 latents get decoded to 13 frames, not 18) + # in order to achieve the intended causal structure... + # anyway, that's why we have to remove some frames here. + # mochi-VAE does the slicing at each TGrow (save compute/mem?) + # but I think it's basically the same + return x[:, 5:] + + def forward(self, x): + return self.c(x) + +@torch.no_grad() +def main(): + """Run TAEM1 roundtrip reconstruction on the given video paths.""" + import sys + import cv2 # no highly esteemed deed is commemorated here + + class VideoTensorReader: + def __init__(self, video_file_path): + self.cap = cv2.VideoCapture(video_file_path) + assert self.cap.isOpened(), f"Could not load {video_file_path}" + self.fps = self.cap.get(cv2.CAP_PROP_FPS) + def __iter__(self): + return self + def __next__(self): + ret, frame = self.cap.read() + if not ret: + self.cap.release() + raise StopIteration # End of video or error + return torch.from_numpy(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).permute(2, 0, 1) # BGR HWC -> RGB CHW + + class VideoTensorWriter: + def __init__(self, video_file_path, width_height, fps=30): + self.writer = cv2.VideoWriter(video_file_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, width_height) + assert self.writer.isOpened(), f"Could not create writer for {video_file_path}" + def write(self, frame_tensor): + assert frame_tensor.ndim == 3 and frame_tensor.shape[0] == 3, f"{frame_tensor.shape}??" + self.writer.write(cv2.cvtColor(frame_tensor.permute(1, 2, 0).numpy(), cv2.COLOR_RGB2BGR)) # RGB CHW -> BGR HWC + def __del__(self): + if hasattr(self, 'writer'): + self.writer.release() + + dev = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu") + dtype = torch.float16 + print("Using device", dev, "and dtype", dtype) + taem1 = TAEM1().to(dev, dtype) + for video_path in sys.argv[1:]: + print(f"Processing {video_path}...") + video_in = VideoTensorReader(video_path) + video = torch.stack(list(video_in), 0)[None] + vid_dev = video.to(dev, dtype).div_(255.0) + # convert to device tensor + if video.numel() < 100_000_000: + print(f" {video_path} seems small enough, will process all frames in parallel") + # convert to device tensor + vid_enc = taem1.encode_video(vid_dev) + print(f" Encoded {video_path} -> {vid_enc.shape}. Decoding...") + vid_dec = taem1.decode_video(vid_enc) + print(f" Decoded {video_path} -> {vid_dec.shape}") + else: + print(f" {video_path} seems large, will process each frame sequentially") + # convert to device tensor + vid_enc = taem1.encode_video(vid_dev, parallel=False) + print(f" Encoded {video_path} -> {vid_enc.shape}. Decoding...") + vid_dec = taem1.decode_video(vid_enc, parallel=False) + print(f" Decoded {video_path} -> {vid_dec.shape}") + video_out_path = video_path + ".reconstructed_by_taem1.mp4" + video_out = VideoTensorWriter(video_out_path, (vid_dec.shape[-1], vid_dec.shape[-2]), fps=int(round(video_in.fps))) + for frame in vid_dec.clamp_(0, 1).mul_(255).round_().byte().cpu()[0]: + video_out.write(frame) + print(f" Saved to {video_out_path}") + +if __name__ == "__main__": + main() diff --git a/modules/ui_video.py b/modules/ui_video.py index 05a3bfc09..905cdfb95 100644 --- a/modules/ui_video.py +++ b/modules/ui_video.py @@ -1,7 +1,7 @@ import gradio as gr from modules import shared, sd_models, timer, images, ui_common, ui_sections, ui_symbols, call_queue, generation_parameters_copypaste from modules.ui_components import ToolButton -from modules.video_models import models_def, video_utils +from modules.video_models import models_def, video_utils, video_load def engine_change(engine): @@ -21,7 +21,7 @@ def model_change(engine, model): sd_models.unload_model_weights() msg = 'Video model unloaded' else: - msg = video_utils.load_model(selected) + msg = video_load.load_model(selected) else: sd_models.unload_model_weights() msg = 'Video model unloaded' @@ -84,7 +84,7 @@ def create_ui(): steps, sampler_index = ui_sections.create_sampler_and_steps_selection(None, "video") with gr.Row(): sampler_shift = gr.Slider(label='Sampler shift', minimum=0.0, maximum=20.0, step=0.1, value=7.0, elem_id="video_scheduler_shift") - dynamic_shift = gr.Checkbox(label='Dynamic shift', value=False, elem_id="video_dynamic_shift", interactive=False) + dynamic_shift = gr.Checkbox(label='Dynamic shift', value=False, elem_id="video_dynamic_shift", interactive=False) # TODO video: dynamic shift with gr.Row(): guidance_scale = gr.Slider(label='Guidance scale', minimum=0.0, maximum=14.0, step=0.1, value=6.0, elem_id="video_guidance_scale") guidance_true = gr.Slider(label='True guidance', minimum=0.0, maximum=14.0, step=0.1, value=1.0, elem_id="video_guidance_true") diff --git a/modules/video_models/models_def.py b/modules/video_models/models_def.py index 82e6f108a..5ff58fd82 100644 --- a/modules/video_models/models_def.py +++ b/modules/video_models/models_def.py @@ -17,6 +17,7 @@ class Model(): te_folder: str = 'text_encoder' te_hijack: bool = True vae_hijack: bool = True + vae_remote: bool = False models = { @@ -25,18 +26,21 @@ models = { Model(name='None'), Model(name='Hunyuan Video T2V', url='https://huggingface.co/tencent/HunyuanVideo', + vae_remote=True, repo='hunyuanvideo-community/HunyuanVideo', repo_cls=diffusers.HunyuanVideoPipeline, te_cls=transformers.LlamaModel, dit_cls=diffusers.HunyuanVideoTransformer3DModel), Model(name='Hunyuan Video I2V', # https://github.com/huggingface/diffusers/pull/10983 url='https://huggingface.co/tencent/HunyuanVideo-I2V', + vae_remote=True, repo='hunyuanvideo-community/HunyuanVideo-I2V', repo_cls=diffusers.HunyuanVideoImageToVideoPipeline, te_cls=transformers.LlavaForConditionalGeneration, dit_cls=diffusers.HunyuanVideoTransformer3DModel), Model(name='SkyReels Hunyuan T2V', # https://github.com/huggingface/diffusers/pull/10837 url='https://huggingface.co/Skywork/SkyReels-V1-Hunyuan-T2V', + vae_remote=True, repo='hunyuanvideo-community/HunyuanVideo', repo_cls=diffusers.HunyuanVideoPipeline, te_cls=transformers.LlamaModel, @@ -45,6 +49,7 @@ models = { dit_cls=diffusers.HunyuanVideoTransformer3DModel), Model(name='SkyReels Hunyuan I2V', # https://github.com/huggingface/diffusers/pull/10837 url='https://huggingface.co/Skywork/SkyReels-V1-Hunyuan-I2V', + vae_remote=True, repo='hunyuanvideo-community/HunyuanVideo', te_cls=transformers.LlamaModel, dit='Skywork/SkyReels-V1-Hunyuan-I2V', @@ -52,6 +57,7 @@ models = { dit_cls=diffusers.HunyuanVideoTransformer3DModel), Model(name='Fast Hunyuan T2V', # https://github.com/hao-ai-lab/FastVideo/blob/8a77cf22c9b9e7f931f42bc4b35d21fd91d24e45/fastvideo/models/hunyuan/inference.py#L213 url='https://huggingface.co/FastVideo/FastHunyuan', + vae_remote=True, repo='hunyuanvideo-community/HunyuanVideo', repo_cls=diffusers.HunyuanVideoPipeline, te_cls=transformers.LlamaModel, @@ -97,6 +103,33 @@ models = { te_cls=transformers.T5EncoderModel, dit_cls=diffusers.LTXVideoTransformer3DModel), ], + 'WAN Video': [ + Model(name='None'), + Model(name='WAN 2.1 1.3B T2V', + url='https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers', + repo='Wan-AI/Wan2.1-T2V-1.3B-Diffusers', + repo_cls=diffusers.WanPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.WanTransformer3DModel), + Model(name='WAN 2.1 14B T2V', + url='https://huggingface.co/Wan-AI/Wan2.1-T2V-14B-Diffusers', + repo='Wan-AI/Wan2.1-T2V-14B-Diffusers', + repo_cls=diffusers.WanPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.WanTransformer3DModel), + Model(name='WAN 2.1 14B I2V 480p', + url='https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P-Diffusers', + repo='Wan-AI/Wan2.1-I2V-14B-480P-Diffusers', + repo_cls=diffusers.WanImageToVideoPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.WanTransformer3DModel), + Model(name='WAN 2.1 14B I2V 720p', + url='https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-720P-Diffusers', + repo='Wan-AI/Wan2.1-I2V-14B-720P-Diffusers', + repo_cls=diffusers.WanImageToVideoPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.WanTransformer3DModel), + ], 'Mochi Video': [ Model(name='None'), Model(name='Mochi 1 T2V', diff --git a/modules/video_models/run_allegro.py b/modules/video_models/run_allegro.py index be1b0da8d..091be7333 100644 --- a/modules/video_models/run_allegro.py +++ b/modules/video_models/run_allegro.py @@ -1,7 +1,7 @@ import os import time from modules import shared, errors, sd_models, processing, devices, images, ui_common -from modules.video_models import models_def, video_utils +from modules.video_models import models_def, video_utils, video_load, video_vae debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -14,7 +14,7 @@ def generate(*args, **kwargs): found = [model.name for model in models_def.models.get(engine, [])] selected: models_def.Model = [m for m in models_def.models[engine] if m.name == model][0] if len(found) > 0 else None if not shared.sd_loaded or 'Allegro' not in shared.sd_model.__class__.__name__: - video_utils.load_model(selected) + video_load.load_model(selected) if not shared.sd_loaded or 'Allegro' not in shared.sd_model.__class__.__name__: return video_utils.queue_err('model not loaded') debug(f'Video generate: task={task_id} args={args} kwargs={kwargs}') @@ -54,13 +54,13 @@ def generate(*args, **kwargs): # set args processing.fix_seed(p) - video_utils.set_vae_params(p.frames, vae_tile_frames) + video_vae.set_vae_params(p) video_utils.set_prompt(p) p.task_args['output_type'] = 'pil' p.ops.append('video') orig_dynamic_shift = shared.opts.schedulers_dynamic_shift orig_sampler_shift = shared.opts.schedulers_shift - shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift # todo video sampler dynamic shift + shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift shared.opts.data['schedulers_shift'] = sampler_shift debug(f'Video: task_args={p.task_args}') diff --git a/modules/video_models/run_cog.py b/modules/video_models/run_cog.py index 5fbe1d320..0aaee4b97 100644 --- a/modules/video_models/run_cog.py +++ b/modules/video_models/run_cog.py @@ -1,7 +1,7 @@ import os import time from modules import shared, errors, sd_models, processing, devices, images, ui_common -from modules.video_models import models_def, video_utils +from modules.video_models import models_def, video_utils, video_load, video_vae debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -14,7 +14,7 @@ def generate(*args, **kwargs): found = [model.name for model in models_def.models.get(engine, [])] selected: models_def.Model = [m for m in models_def.models[engine] if m.name == model][0] if len(found) > 0 else None if not shared.sd_loaded or 'Cog' not in shared.sd_model.__class__.__name__: - video_utils.load_model(selected) + video_load.load_model(selected) if not shared.sd_loaded or 'Cog' not in shared.sd_model.__class__.__name__: return video_utils.queue_err('model not loaded') debug(f'Video generate: task={task_id} args={args} kwargs={kwargs}') @@ -54,13 +54,13 @@ def generate(*args, **kwargs): # set args processing.fix_seed(p) - video_utils.set_vae_params(p.frames, vae_tile_frames) + video_vae.set_vae_params(p) video_utils.set_prompt(p) p.task_args['output_type'] = 'pil' p.ops.append('video') orig_dynamic_shift = shared.opts.schedulers_dynamic_shift orig_sampler_shift = shared.opts.schedulers_shift - shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift # todo video sampler dynamic shift + shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift shared.opts.data['schedulers_shift'] = sampler_shift debug(f'Video: task_args={p.task_args}') diff --git a/modules/video_models/run_hunyuan.py b/modules/video_models/run_hunyuan.py index aa076f935..0f722b890 100644 --- a/modules/video_models/run_hunyuan.py +++ b/modules/video_models/run_hunyuan.py @@ -1,7 +1,7 @@ import os import time from modules import shared, errors, sd_models, processing, devices, images, ui_common -from modules.video_models import models_def, video_utils +from modules.video_models import models_def, video_utils, video_load, video_vae debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -14,7 +14,7 @@ def generate(*args, **kwargs): found = [model.name for model in models_def.models.get(engine, [])] selected: models_def.Model = [m for m in models_def.models[engine] if m.name == model][0] if len(found) > 0 else None if not shared.sd_loaded or 'Hunyuan' not in shared.sd_model.__class__.__name__: - video_utils.load_model(selected) + video_load.load_model(selected) if not shared.sd_loaded or 'Hunyuan' not in shared.sd_model.__class__.__name__: return video_utils.queue_err('model not loaded') debug(f'Video generate: task={task_id} args={args} kwargs={kwargs}') @@ -38,6 +38,9 @@ def generate(*args, **kwargs): vae_tile_frames=int(vae_tile_frames), override_settings=override_settings, ) + if p.vae_type == 'Remote' and not selected.vae_remote: + shared.log.warning(f'Video: model={selected.name} remote vae not supported') + p.vae_type = 'Default' p.scripts = None p.script_args = None p.state = ui_state @@ -54,13 +57,13 @@ def generate(*args, **kwargs): # set args processing.fix_seed(p) - video_utils.set_vae_params(p.frames, vae_tile_frames) + video_vae.set_vae_params(p) video_utils.set_prompt(p) - p.task_args['output_type'] = 'pil' + p.task_args['output_type'] = 'latent' if (p.vae_type == 'Remote') else 'pil' p.ops.append('video') orig_dynamic_shift = shared.opts.schedulers_dynamic_shift orig_sampler_shift = shared.opts.schedulers_shift - shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift # todo video sampler dynamic shift + shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift shared.opts.data['schedulers_shift'] = sampler_shift debug(f'Video: task_args={p.task_args}') diff --git a/modules/video_models/run_ltx.py b/modules/video_models/run_ltx.py index ba01e45a4..ef3fa08b8 100644 --- a/modules/video_models/run_ltx.py +++ b/modules/video_models/run_ltx.py @@ -1,7 +1,7 @@ import os import time from modules import shared, errors, sd_models, processing, devices, images, ui_common -from modules.video_models import models_def, video_utils +from modules.video_models import models_def, video_utils, video_load, video_vae debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -14,7 +14,7 @@ def generate(*args, **kwargs): found = [model.name for model in models_def.models.get(engine, [])] selected: models_def.Model = [m for m in models_def.models[engine] if m.name == model][0] if len(found) > 0 else None if not shared.sd_loaded or 'LTX' not in shared.sd_model.__class__.__name__: - video_utils.load_model(selected) + video_load.load_model(selected) if not shared.sd_loaded or 'LTX' not in shared.sd_model.__class__.__name__: return video_utils.queue_err('model not loaded') debug(f'Video generate: task={task_id} args={args} kwargs={kwargs}') @@ -54,13 +54,13 @@ def generate(*args, **kwargs): # set args processing.fix_seed(p) - video_utils.set_vae_params(p.frames, vae_tile_frames) + video_vae.set_vae_params(p) video_utils.set_prompt(p) p.task_args['output_type'] = 'pil' p.ops.append('video') orig_dynamic_shift = shared.opts.schedulers_dynamic_shift orig_sampler_shift = shared.opts.schedulers_shift - shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift # todo video sampler dynamic shift + shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift shared.opts.data['schedulers_shift'] = sampler_shift debug(f'Video: task_args={p.task_args}') diff --git a/modules/video_models/run_mochi.py b/modules/video_models/run_mochi.py index fff3a6d04..ef705880e 100644 --- a/modules/video_models/run_mochi.py +++ b/modules/video_models/run_mochi.py @@ -1,7 +1,7 @@ import os import time from modules import shared, errors, sd_models, processing, devices, images, ui_common -from modules.video_models import models_def, video_utils +from modules.video_models import models_def, video_utils, video_load, video_vae debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -14,7 +14,7 @@ def generate(*args, **kwargs): found = [model.name for model in models_def.models.get(engine, [])] selected: models_def.Model = [m for m in models_def.models[engine] if m.name == model][0] if len(found) > 0 else None if not shared.sd_loaded or 'Mochi' not in shared.sd_model.__class__.__name__: - video_utils.load_model(selected) + video_load.load_model(selected) if not shared.sd_loaded or 'Mochi' not in shared.sd_model.__class__.__name__: return video_utils.queue_err('model not loaded') debug(f'Video generate: task={task_id} args={args} kwargs={kwargs}') @@ -54,13 +54,13 @@ def generate(*args, **kwargs): # set args processing.fix_seed(p) - video_utils.set_vae_params(p.frames, vae_tile_frames) + video_vae.set_vae_params(p) video_utils.set_prompt(p) p.task_args['output_type'] = 'pil' p.ops.append('video') orig_dynamic_shift = shared.opts.schedulers_dynamic_shift orig_sampler_shift = shared.opts.schedulers_shift - shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift # todo video sampler dynamic shift + shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift shared.opts.data['schedulers_shift'] = sampler_shift debug(f'Video: task_args={p.task_args}') diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py new file mode 100644 index 000000000..2c8634f1e --- /dev/null +++ b/modules/video_models/video_load.py @@ -0,0 +1,80 @@ +import os +import time +from modules import shared, errors, sd_models, sd_checkpoint, model_quant, devices +from modules.video_models import models_def, video_utils, video_vae + + +loaded_model = None +debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def load_model(selected: models_def.Model): + if selected is None: + return + global loaded_model # pylint: disable=global-statement + if loaded_model == selected.name: + return + sd_models.unload_model_weights() + t0 = time.time() + + # text encoder + try: + quant_args = model_quant.create_config(module='Text Encoder') + debug(f'Video load: module=te repo="{selected.te or selected.repo}" folder="{selected.te_folder}" cls={selected.te_cls.__name__} quant={video_utils.get_quant(quant_args)}') + text_encoder = selected.te_cls.from_pretrained( + pretrained_model_name_or_path=selected.te or selected.repo, + subfolder=selected.te_folder, + cache_dir=shared.opts.hfcache_dir, + torch_dtype=devices.dtype, + **quant_args + ) + except Exception as e: + shared.log.error(f'video load: module=te cls={selected.te_cls.__name__} {e}') + errors.display(e, 'video') + text_encoder = None + + # transformer + try: + quant_args = model_quant.create_config(module='Model') + debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" folder="{selected.dit_folder}" cls={selected.dit_cls.__name__} quant={video_utils.get_quant(quant_args)}') + transformer = selected.dit_cls.from_pretrained( + pretrained_model_name_or_path=selected.dit or selected.repo, + subfolder=selected.dit_folder, + torch_dtype=devices.dtype, + cache_dir=shared.opts.hfcache_dir, + **quant_args + ) + except Exception as e: + shared.log.error(f'video load: module=transformer cls={selected.dit_cls.__name__} {e}') + errors.display(e, 'video') + transformer = None + + # model + try: + debug(f'Video load: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__}') + shared.sd_model = selected.repo_cls.from_pretrained( + pretrained_model_name_or_path=selected.repo, + transformer=transformer, + text_encoder=text_encoder, + cache_dir=shared.opts.hfcache_dir, + torch_dtype=devices.dtype, + ) + except Exception as e: + shared.log.error(f'video load: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__} {e}') + errors.display(e, 'video') + + t1 = time.time() + shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(selected.repo) + shared.sd_model.sd_model_hash = None + sd_models.set_diffuser_options(shared.sd_model) + if selected.vae_hijack: + shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode + shared.sd_model.vae.decode = video_vae.hijack_vae_decode + if selected.te_hijack: + shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt + shared.sd_model.encode_prompt = video_utils.hijack_encode_prompt + shared.sd_model.vae.enable_slicing() + loaded_model = selected.name + msg = f'Video load: cls={shared.sd_model.__class__.__name__} model="{selected.name}" time={t1-t0:.2f}' + shared.log.info(msg) + return msg diff --git a/modules/video_models/video_utils.py b/modules/video_models/video_utils.py index c2b3878cf..0ed4791eb 100644 --- a/modules/video_models/video_utils.py +++ b/modules/video_models/video_utils.py @@ -1,7 +1,6 @@ import os import time -from modules import shared, errors, timer, sd_models, sd_checkpoint, model_quant, devices -from modules.video_models import models_def +from modules import shared, sd_models, timer debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -29,31 +28,6 @@ def set_prompt(p): p.task_args['negative_prompt'] = p.negative_prompt -def set_vae_params(frames, tile_frames): - if tile_frames > frames: - if hasattr(shared.sd_model.vae, 'tile_sample_min_num_frames'): - shared.sd_model.vae.tile_sample_min_num_frames = tile_frames - if hasattr(shared.sd_model.vae, 'use_framewise_decoding'): - shared.sd_model.vae.use_framewise_decoding = True - if hasattr(shared.sd_model.vae, 'enable_tiling'): - shared.sd_model.vae.enable_tiling() - 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() - - -def hijack_vae_decode(*args, **kwargs): - t0 = time.time() - shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) - res = shared.sd_model.vae.orig_decode(*args, **kwargs) - t1 = time.time() - timer.process.add('vae', t1-t0) - debug(f'Video decode: vae={shared.sd_model.vae.__class__.__name__} time={t1-t0:.2f}') - return res - - def hijack_encode_prompt(*args, **kwargs): t0 = time.time() res = shared.sd_model.orig_encode_prompt(*args, **kwargs) @@ -62,77 +36,3 @@ def hijack_encode_prompt(*args, **kwargs): debug(f'Video encode: te={shared.sd_model.text_encoder.__class__.__name__} time={t1-t0:.2f}') shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) return res - -loaded_model = None - - -def load_model(selected: models_def.Model): - if selected is None: - return - global loaded_model # pylint: disable=global-statement - if loaded_model == selected.name: - return - sd_models.unload_model_weights() - t0 = time.time() - - # text encoder - try: - quant_args = model_quant.create_config(module='Text Encoder') - debug(f'Video load: module=te repo="{selected.te or selected.repo}" folder="{selected.te_folder}" cls={selected.te_cls.__name__} quant={get_quant(quant_args)}') - text_encoder = selected.te_cls.from_pretrained( - pretrained_model_name_or_path=selected.te or selected.repo, - subfolder=selected.te_folder, - cache_dir=shared.opts.hfcache_dir, - torch_dtype=devices.dtype, - **quant_args - ) - except Exception as e: - shared.log.error(f'video load: module=te cls={selected.te_cls.__name__} {e}') - errors.display(e, 'video') - text_encoder = None - - # transformer - try: - quant_args = model_quant.create_config(module='Model') - debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" folder="{selected.dit_folder}" cls={selected.dit_cls.__name__} quant={get_quant(quant_args)}') - transformer = selected.dit_cls.from_pretrained( - pretrained_model_name_or_path=selected.dit or selected.repo, - subfolder=selected.dit_folder, - torch_dtype=devices.dtype, - cache_dir=shared.opts.hfcache_dir, - **quant_args - ) - except Exception as e: - shared.log.error(f'video load: module=transformer cls={selected.dit_cls.__name__} {e}') - errors.display(e, 'video') - transformer = None - - # model - try: - debug(f'Video load: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__}') - shared.sd_model = selected.repo_cls.from_pretrained( - pretrained_model_name_or_path=selected.repo, - transformer=transformer, - text_encoder=text_encoder, - cache_dir=shared.opts.hfcache_dir, - torch_dtype=devices.dtype, - ) - except Exception as e: - shared.log.error(f'video load: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__} {e}') - errors.display(e, 'video') - - t1 = time.time() - sd_models.set_diffuser_options(shared.sd_model) - shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(selected.repo) - shared.sd_model.sd_model_hash = None - if selected.vae_hijack: - shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode - shared.sd_model.vae.decode = hijack_vae_decode - if selected.te_hijack: - shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt - shared.sd_model.encode_prompt = hijack_encode_prompt - shared.sd_model.vae.enable_slicing() - loaded_model = selected.name - msg = f'Video load: cls={shared.sd_model.__class__.__name__} model="{selected.name}" time={t1-t0:.2f}' - shared.log.info(msg) - return msg diff --git a/modules/video_models/video_vae.py b/modules/video_models/video_vae.py new file mode 100644 index 000000000..bd552f66c --- /dev/null +++ b/modules/video_models/video_vae.py @@ -0,0 +1,62 @@ +import os +import time +from modules import shared, sd_models, devices, timer + + +debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None +vae_type = None + + +def set_vae_params(p): + global vae_type # pylint: disable=global-statement + vae_type = p.vae_type + if p.vae_tile_frames > p.frames: + if hasattr(shared.sd_model.vae, 'tile_sample_min_num_frames'): + shared.sd_model.vae.tile_sample_min_num_frames = p.vae_tile_frames + if hasattr(shared.sd_model.vae, 'use_framewise_decoding'): + shared.sd_model.vae.use_framewise_decoding = True + if hasattr(shared.sd_model.vae, 'enable_tiling'): + shared.sd_model.vae.enable_tiling() + 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() + + +def vae_decode_tiny(latents): + if 'Hunyuan' in shared.sd_model.__class__.__name__: + variant = 'TAE HunyuanVideo' + elif 'Mochi' in shared.sd_model.__class__.__name__: + variant = 'TAE MochiVideo' + elif 'WAN' in shared.sd_model.__class__.__name__: + variant = 'TAE WanVideo' + else: + shared.log.warning(f'Video VAE: type=Tiny cls={shared.sd_model.__class__.__name__} not supported') + return None + from modules import sd_vae_taesd + vae = sd_vae_taesd.get_model(variant) + if vae is None: + return None + debug(f'Video VAE: 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 = images.transpose(1, 2).mul_(2).sub_(1) + return (images, None) + + +def hijack_vae_decode(*args, **kwargs): + t0 = time.time() + res = None + if vae_type == 'Tiny': + res = vae_decode_tiny(args[0]) + if vae_type == 'Remote': + pass + if res is None: + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) + res = shared.sd_model.vae.orig_decode(*args, **kwargs) + t1 = time.time() + timer.process.add('vae', t1-t0) + debug(f'Video decode: type={vae_type} vae={shared.sd_model.vae.__class__.__name__} latents={args[0].shape} time={t1-t0:.2f}') + return res diff --git a/scripts/cogvideo.py b/scripts/cogvideo.py index 7184dd946..c18b4eb2f 100644 --- a/scripts/cogvideo.py +++ b/scripts/cogvideo.py @@ -50,7 +50,7 @@ class Script(scripts.Script): return [model, sampler, frames, guidance, offload, override, video_type, duration, loop, pad, interpolate, image, video] def load(self, model): - if (shared.sd_model_type != 'cogvideox' or shared.sd_model.sd_model_checkpoint != model) and model != 'None': + if (shared.sd_model_type != 'cogvideo' or shared.sd_model.sd_model_checkpoint != model) and model != 'None': sd_models.unload_model_weights('model') shared.log.info(f'CogVideoX load: model="{model}"') try: @@ -64,7 +64,7 @@ class Script(scripts.Script): shared.log.error(f'Load CogVideoX: {e}') if debug: errors.display(e, 'CogVideoX') - if shared.sd_model_type == 'cogvideox' and model != 'None': + if shared.sd_model_type == 'cogvideo' and model != 'None': shared.sd_model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m', ncols=80, colour='#327fba') shared.log.debug(f'CogVideoX load: class="{shared.sd_model.__class__.__name__}"') if shared.sd_model is not None and model == 'None': @@ -74,7 +74,7 @@ class Script(scripts.Script): devices.torch_gc() def offload(self, offload): - if shared.sd_model_type != 'cogvideox': + if shared.sd_model_type != 'cogvideo': return if offload == 'none': sd_models.move_model(shared.sd_model, devices.device) @@ -131,7 +131,7 @@ class Script(scripts.Script): return img def generate(self, p: processing.StableDiffusionProcessing, model: str): - if shared.sd_model_type != 'cogvideox': + if shared.sd_model_type != 'cogvideo': return [] shared.log.info(f'CogVideoX: sampler={p.sampler} steps={p.steps} frames={p.frames} width={p.width} height={p.height} seed={p.seed} guidance={p.guidance}') if p.sampler == 'DDIM':