mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 16:54:33 +02:00
explicit seedvr implementation
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
+5
-2
@@ -1,6 +1,6 @@
|
||||
# Change Log for SD.Next
|
||||
|
||||
## Update for 2026-07-16
|
||||
## Update for 2026-07-18
|
||||
|
||||
- **Compute**
|
||||
- torch: update to `2.13.0` for CUDA, ROCm, IPEX
|
||||
@@ -10,9 +10,12 @@
|
||||
- sdnq quantization optimizations
|
||||
- sdnq attention optimizations
|
||||
- sdnq separate dit/te settings
|
||||
- **Features**
|
||||
- SeedVR enhanced support
|
||||
- **Fixes**
|
||||
- upscaler auto-refresh to catch chainner upscalers that are not loaded on first attempt
|
||||
- lora loader support diffusers trainer
|
||||
- lora support diffusers trainer
|
||||
- flux1 load t5
|
||||
|
||||
## Update for 2026-07-14
|
||||
|
||||
|
||||
Submodule extensions-builtin/sdnext-modernui updated: 07890f0d7b...b987359783
@@ -28,6 +28,8 @@ class UpscalerSeedVR(Upscaler):
|
||||
]
|
||||
self.model = None
|
||||
self.model_loaded = None
|
||||
self.tile_size = 1024
|
||||
self.tile_overlap = 0.25
|
||||
self.device = devices.device
|
||||
|
||||
def load_model(self, path: str):
|
||||
@@ -60,8 +62,9 @@ class UpscalerSeedVR(Upscaler):
|
||||
}
|
||||
t1 = time.time()
|
||||
self.model.dit.config = self.model.config.dit
|
||||
self.model.vae.tile_sample_min_size = 1024
|
||||
self.model.vae.tile_latent_min_size = 128
|
||||
self.model.vae.tile_sample_min_size = self.tile_size
|
||||
self.model.vae.tile_latent_min_size = self.tile_size // 8
|
||||
self.model.vae.tile_overlap_factor = self.tile_overlap
|
||||
|
||||
self.model = do_post_load_quant(self.model, allow=True)
|
||||
|
||||
@@ -136,26 +139,34 @@ class UpscalerSeedVR(Upscaler):
|
||||
devices.torch_gc()
|
||||
return result
|
||||
|
||||
def do_upscale(self, img: Image.Image, selected_file):
|
||||
def do_upscale(self, img: Image.Image, selected_file, cfg_scale: float = 3.5, cfg_rescale: float = 0.0, steps: int = 1, seed: int = -1, scale: float | None = None, tile_size: int = 1024, tile_overlap: float = 0.25):
|
||||
self.load_model(selected_file)
|
||||
if self.model is None:
|
||||
return img
|
||||
|
||||
from modules.seedvr.src.core import generation
|
||||
|
||||
self.scale = self.scale if scale is None else scale
|
||||
self.tile_size = tile_size if tile_size is not None else self.tile_size
|
||||
self.tile_overlap = tile_overlap if tile_overlap is not None else self.tile_overlap
|
||||
self.model.vae.tile_sample_min_size = self.tile_size
|
||||
self.model.vae.tile_latent_min_size = self.tile_size // 8
|
||||
self.model.vae.tile_overlap_factor = self.tile_overlap
|
||||
width = int(self.scale * img.width) // 8 * 8
|
||||
image_tensor = np.array(img)
|
||||
image_tensor = torch.from_numpy(image_tensor).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) / 255.0
|
||||
|
||||
random.seed()
|
||||
seed = int(random.randrange(4294967294))
|
||||
seed = int(random.randrange(4294967294)) if seed == -1 else int(seed)
|
||||
|
||||
t0 = time.time()
|
||||
with devices.inference_context():
|
||||
result_tensor = generation.generation_loop(
|
||||
runner=self.model,
|
||||
images=image_tensor,
|
||||
cfg_scale=opts.seedvr_cfg_scale,
|
||||
cfg_scale=cfg_scale,
|
||||
cfg_rescale=cfg_rescale,
|
||||
steps=steps,
|
||||
seed=seed,
|
||||
res_w=width,
|
||||
batch_size=1,
|
||||
@@ -163,7 +174,8 @@ class UpscalerSeedVR(Upscaler):
|
||||
device=devices.device,
|
||||
)
|
||||
t1 = time.time()
|
||||
log.info(f'Upscaler: type="{self.name}" model="{selected_file}" scale={self.scale} cfg={opts.seedvr_cfg_scale} seed={seed} time={t1 - t0:.2f}')
|
||||
tiles = getattr(self.model.vae, "tiles", None)
|
||||
log.info(f'Upscaler: type="{self.name}" model="{selected_file}" scale={self.scale} cfg={cfg_scale} seed={seed} tiles={tiles} time={t1 - t0:.2f}')
|
||||
img = convert.to_pil(result_tensor.squeeze())
|
||||
|
||||
if opts.upscaler_unload:
|
||||
|
||||
@@ -18,16 +18,10 @@ Euler ODE solver.
|
||||
"""
|
||||
|
||||
from typing import Callable
|
||||
import itertools
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch.nn import functional as F
|
||||
|
||||
#from ....models.dit_v2 import na
|
||||
|
||||
from ..types import PredictionType
|
||||
from ..utils import expand_dims
|
||||
from .base import Sampler, SamplerModelArgs
|
||||
import itertools
|
||||
|
||||
|
||||
class EulerSampler(Sampler):
|
||||
|
||||
@@ -106,7 +106,7 @@ def cut_videos(videos):
|
||||
return result
|
||||
|
||||
|
||||
def generation_loop(runner, images, cfg_scale=1.0, seed=666, res_w=720, batch_size=90, temporal_overlap=0, progress_callback=None, device:str='cpu'):
|
||||
def generation_loop(runner, images, cfg_scale=1.0, cfg_rescale=0.0, steps=1, seed=666, res_w=720, batch_size=90, temporal_overlap=0, progress_callback=None, device:str='cpu'):
|
||||
"""
|
||||
Main generation loop with context-aware temporal processing
|
||||
|
||||
@@ -137,9 +137,9 @@ def generation_loop(runner, images, cfg_scale=1.0, seed=666, res_w=720, batch_si
|
||||
|
||||
# Configure classifier-free guidance
|
||||
runner.config.diffusion.cfg.scale = cfg_scale
|
||||
runner.config.diffusion.cfg.rescale = 0.0
|
||||
runner.config.diffusion.cfg.rescale = cfg_rescale
|
||||
# Configure sampling steps
|
||||
runner.config.diffusion.timesteps.sampling.steps = 1
|
||||
runner.config.diffusion.timesteps.sampling.steps = steps
|
||||
runner.configure_diffusion()
|
||||
|
||||
# Set random seed
|
||||
|
||||
@@ -309,10 +309,7 @@ class SeedVRPipeline():
|
||||
timestep=args.t.repeat(batch_size),
|
||||
).vid_sample,
|
||||
scale=(
|
||||
cfg_scale
|
||||
if (args.i + 1) / len(self.sampler.timesteps)
|
||||
<= self.config.diffusion.cfg.get("partial", 1)
|
||||
else 1.0
|
||||
cfg_scale if (args.i + 1) / len(self.sampler.timesteps) <= self.config.diffusion.cfg.get("partial", 1) else 1.0
|
||||
),
|
||||
rescale=self.config.diffusion.cfg.rescale,
|
||||
),
|
||||
|
||||
@@ -50,5 +50,6 @@ class SideResize:
|
||||
size = min(width, height)
|
||||
else:
|
||||
size = self.size
|
||||
|
||||
return TVF.resize(image, size, self.interpolation)
|
||||
size_w = int(size) // 8 * 8
|
||||
size_h = int(size * height / width) // 8 * 8
|
||||
return TVF.resize(image, (size_h, size_w), self.interpolation)
|
||||
|
||||
@@ -42,6 +42,8 @@ class PatchIn(nn.Module):
|
||||
) -> torch.Tensor:
|
||||
t, h, w = self.patch_size
|
||||
vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w)
|
||||
if vid.dtype != self.proj.weight.dtype:
|
||||
vid = vid.to(self.proj.weight.dtype)
|
||||
vid = self.proj(vid)
|
||||
return vid
|
||||
|
||||
@@ -63,30 +65,40 @@ class PatchOut(nn.Module):
|
||||
vid: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
t, h, w = self.patch_size
|
||||
if vid.dtype != self.proj.weight.dtype:
|
||||
vid = vid.to(self.proj.weight.dtype)
|
||||
vid = self.proj(vid)
|
||||
vid = rearrange(vid, "b T H W (t h w c) -> b c (T t) (H h) (W w)", t=t, h=h, w=w)
|
||||
return vid
|
||||
|
||||
|
||||
class NaPatchIn(PatchIn):
|
||||
def forward(
|
||||
def forward( # pylint: disable=arguments-differ
|
||||
self,
|
||||
vid: torch.Tensor, # l c
|
||||
vid_shape: torch.LongTensor,
|
||||
) -> torch.Tensor:
|
||||
t, h, w = self.patch_size
|
||||
if not (t == h == w == 1):
|
||||
if not t == h == w == 1:
|
||||
vid, vid_shape = na.rearrange(
|
||||
vid, vid_shape, "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w
|
||||
)
|
||||
for i in range(len(vid)):
|
||||
if h > 1 and vid_shape[i, 1] % h != 0:
|
||||
vid[i] = torch.cat([vid[i][:, :1]] * (h - vid[i].size(1) % h) + [vid[i]], dim=1)
|
||||
if w > 1 and vid_shape[i, 2] % w != 0:
|
||||
vid[i] = torch.cat([vid[i][:, :, :1]] * (w - vid[i].size(2) % w) + [vid[i]], dim=2)
|
||||
vid, vid_shape = na.flatten(vid)
|
||||
# slice vid after patching in when using sequence parallelism
|
||||
vid = slice_inputs(vid, dim=0)
|
||||
if vid.dtype != self.proj.weight.dtype:
|
||||
vid = vid.to(self.proj.weight.dtype)
|
||||
vid = self.proj(vid)
|
||||
return vid, vid_shape
|
||||
|
||||
|
||||
class NaPatchOut(PatchOut):
|
||||
def forward(
|
||||
def forward( # pylint: disable=arguments-differ
|
||||
self,
|
||||
vid: torch.FloatTensor, # l c
|
||||
vid_shape: torch.LongTensor,
|
||||
@@ -96,8 +108,10 @@ class NaPatchOut(PatchOut):
|
||||
torch.LongTensor,
|
||||
]:
|
||||
t, h, w = self.patch_size
|
||||
if vid.dtype != self.proj.weight.dtype:
|
||||
vid = vid.to(self.proj.weight.dtype)
|
||||
vid = self.proj(vid)
|
||||
# gather vid before patching out when enabling sequence parallelism
|
||||
# gather vid before patchting out when enabling sequence parallelism
|
||||
vid = gather_outputs(
|
||||
vid,
|
||||
gather_dim=0,
|
||||
@@ -105,8 +119,13 @@ class NaPatchOut(PatchOut):
|
||||
unpad_shape=vid_shape,
|
||||
cache=cache.namespace("vid"),
|
||||
)
|
||||
if not (t == h == w == 1):
|
||||
if not t == h == w == 1:
|
||||
vid, vid_shape = na.rearrange(
|
||||
vid, vid_shape, "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w
|
||||
)
|
||||
for i in range(len(vid)):
|
||||
if h > 1 and vid_shape[i, 1] % h != 0:
|
||||
vid[i] = vid[i][:, (h - vid_shape[i, 1] % h) :]
|
||||
if w > 1 and vid_shape[i, 2] % w != 0:
|
||||
vid[i] = vid[i][:, :, (w - vid_shape[i, 2] % w) :]
|
||||
return vid, vid_shape
|
||||
|
||||
@@ -45,6 +45,8 @@ class PatchIn(nn.Module):
|
||||
assert vid.size(2) % t == 1
|
||||
vid = torch.cat([vid[:, :, :1]] * (t - 1) + [vid], dim=2)
|
||||
vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w)
|
||||
if vid.dtype != self.proj.weight.dtype:
|
||||
vid = vid.to(self.proj.weight.dtype)
|
||||
vid = self.proj(vid)
|
||||
return vid
|
||||
|
||||
@@ -83,16 +85,22 @@ class NaPatchIn(PatchIn):
|
||||
cache = cache.namespace("patch")
|
||||
vid_shape_before_patchify = cache("vid_shape_before_patchify", lambda: vid_shape)
|
||||
t, h, w = self.patch_size
|
||||
if not (t == h == w == 1):
|
||||
if not t == h == w == 1:
|
||||
vid = na.unflatten(vid, vid_shape)
|
||||
for i in range(len(vid)):
|
||||
if t > 1 and vid_shape_before_patchify[i, 0] % t != 0:
|
||||
vid[i] = torch.cat([vid[i][:1]] * (t - vid[i].size(0) % t) + [vid[i]], dim=0)
|
||||
if h > 1 and vid_shape_before_patchify[i, 1] % h != 0:
|
||||
vid[i] = torch.cat([vid[i][:, :1]] * (h - vid[i].size(1) % h) + [vid[i]], dim=1)
|
||||
if w > 1 and vid_shape_before_patchify[i, 2] % w != 0:
|
||||
vid[i] = torch.cat([vid[i][:, :, :1]] * (w - vid[i].size(2) % w) + [vid[i]], dim=2)
|
||||
vid[i] = rearrange(vid[i], "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w)
|
||||
vid, vid_shape = na.flatten(vid)
|
||||
|
||||
# slice vid after patching in when using sequence parallelism
|
||||
vid = slice_inputs(vid, dim=0)
|
||||
if vid.dtype != self.proj.weight.dtype:
|
||||
vid = vid.to(self.proj.weight.dtype)
|
||||
vid = self.proj(vid)
|
||||
return vid, vid_shape
|
||||
|
||||
@@ -111,17 +119,23 @@ class NaPatchOut(PatchOut):
|
||||
vid_shape_before_patchify = cache.get("vid_shape_before_patchify")
|
||||
|
||||
t, h, w = self.patch_size
|
||||
if vid.dtype != self.proj.weight.dtype:
|
||||
vid = vid.to(self.proj.weight.dtype)
|
||||
vid = self.proj(vid)
|
||||
# gather vid before patching out when enabling sequence parallelism
|
||||
vid = gather_outputs(
|
||||
vid, gather_dim=0, padding_dim=0, unpad_shape=vid_shape, cache=cache.namespace("vid")
|
||||
)
|
||||
if not (t == h == w == 1):
|
||||
if not t == h == w == 1:
|
||||
vid = na.unflatten(vid, vid_shape)
|
||||
for i in range(len(vid)):
|
||||
vid[i] = rearrange(vid[i], "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w)
|
||||
if t > 1 and vid_shape_before_patchify[i, 0] % t != 0:
|
||||
vid[i] = vid[i][(t - vid_shape_before_patchify[i, 0] % t) :]
|
||||
if h > 1 and vid_shape_before_patchify[i, 1] % h != 0:
|
||||
vid[i] = vid[i][:, (h - vid_shape_before_patchify[i, 1] % h) :]
|
||||
if w > 1 and vid_shape_before_patchify[i, 2] % w != 0:
|
||||
vid[i] = vid[i][:, :, (w - vid_shape_before_patchify[i, 2] % w) :]
|
||||
vid, vid_shape = na.flatten(vid)
|
||||
|
||||
return vid, vid_shape
|
||||
|
||||
@@ -1044,7 +1044,7 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
|
||||
norm_num_groups: int = 32,
|
||||
sample_size: int = 32,
|
||||
scaling_factor: float = 0.18215,
|
||||
force_upcast: float = True,
|
||||
force_upcast: float = False,
|
||||
attention: bool = True,
|
||||
temporal_scale_num: int = 0,
|
||||
slicing_up_num: int = 0,
|
||||
@@ -1244,12 +1244,14 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
|
||||
blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
|
||||
row_limit = self.tile_latent_min_size - blend_extent
|
||||
rows = []
|
||||
self.tiles = 0
|
||||
for i in range(0, x.shape[3], overlap_size):
|
||||
row = []
|
||||
for j in range(0, x.shape[4], overlap_size):
|
||||
tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
|
||||
tile = self._encode(tile)
|
||||
row.append(tile)
|
||||
self.tiles += 1
|
||||
rows.append(row)
|
||||
result_rows = []
|
||||
for i, row in enumerate(rows):
|
||||
@@ -1329,6 +1331,7 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
|
||||
self.spatial_downsample_factor = spatial_downsample_factor
|
||||
self.temporal_downsample_factor = temporal_downsample_factor
|
||||
self.freeze_encoder = freeze_encoder
|
||||
self.freeze_encoder = True
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def forward(self, x: torch.FloatTensor) -> CausalAutoencoderOutput:
|
||||
|
||||
@@ -1,936 +0,0 @@
|
||||
# Copyright (c) 2023 HuggingFace Team
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
|
||||
# SPDX-License-Identifier: Apache License, Version 2.0 (the "License")
|
||||
#
|
||||
# This file has been modified by ByteDance Ltd. and/or its affiliates. on 1st June 2025
|
||||
#
|
||||
# Original file was released under Apache License, Version 2.0 (the "License"), with the full license text
|
||||
# available at http://www.apache.org/licenses/LICENSE-2.0.
|
||||
#
|
||||
# This modified file is released under the same license.
|
||||
|
||||
from contextlib import nullcontext
|
||||
from typing import Optional, Tuple, Literal, Callable, Union
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
|
||||
from einops import rearrange
|
||||
from ....common.half_precision_fixes import safe_pad_operation
|
||||
from ....common.logger import get_logger
|
||||
from .causal_inflation_lib import InflatedCausalConv3d, causal_norm_wrapper, init_causal_conv3d, remove_head
|
||||
from .context_parallel_lib import causal_conv_gather_outputs, causal_conv_slice_inputs
|
||||
from .global_config import set_norm_limit
|
||||
from .types import CausalAutoencoderOutput, CausalDecoderOutput, CausalEncoderOutput, MemoryState, _inflation_mode_t, _memory_device_t, _receptive_field_t, _selective_checkpointing_t
|
||||
|
||||
|
||||
logger = get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
# Fake func, no checkpointing is required for inference
|
||||
def gradient_checkpointing(module: Union[Callable, nn.Module], *args, enabled: bool, **kwargs):
|
||||
return module(*args, **kwargs)
|
||||
|
||||
class ResnetBlock2D(nn.Module):
|
||||
r"""
|
||||
A Resnet block.
|
||||
|
||||
Parameters:
|
||||
in_channels (`int`): The number of channels in the input.
|
||||
out_channels (`int`, *optional*, default to be `None`):
|
||||
The number of output channels for the first conv2d layer.
|
||||
If None, same as `in_channels`.
|
||||
dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, *, in_channels: int, out_channels: Optional[int] = None, dropout: float = 0.0
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.nonlinearity = nn.SiLU()
|
||||
|
||||
self.norm1 = torch.nn.GroupNorm(
|
||||
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
|
||||
)
|
||||
|
||||
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.norm2 = torch.nn.GroupNorm(
|
||||
num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
|
||||
)
|
||||
|
||||
self.dropout = torch.nn.Dropout(dropout)
|
||||
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.use_in_shortcut = self.in_channels != out_channels
|
||||
|
||||
self.conv_shortcut = None
|
||||
if self.use_in_shortcut:
|
||||
self.conv_shortcut = nn.Conv2d(
|
||||
in_channels, out_channels, kernel_size=1, stride=1, padding=0
|
||||
)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
|
||||
hidden = input_tensor
|
||||
|
||||
hidden = self.norm1(hidden)
|
||||
hidden = self.nonlinearity(hidden)
|
||||
hidden = self.conv1(hidden)
|
||||
|
||||
hidden = self.norm2(hidden)
|
||||
hidden = self.nonlinearity(hidden)
|
||||
hidden = self.dropout(hidden)
|
||||
hidden = self.conv2(hidden)
|
||||
|
||||
if self.conv_shortcut is not None:
|
||||
input_tensor = self.conv_shortcut(input_tensor)
|
||||
|
||||
output_tensor = input_tensor + hidden
|
||||
|
||||
return output_tensor
|
||||
|
||||
class Upsample3D(nn.Module):
|
||||
"""A 3D upsampling layer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
inflation_mode: _inflation_mode_t = "tail",
|
||||
temporal_up: bool = False,
|
||||
spatial_up: bool = True,
|
||||
slicing: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.conv = init_causal_conv3d(
|
||||
self.channels, self.channels, kernel_size=3, padding=1, inflation_mode=inflation_mode
|
||||
)
|
||||
|
||||
self.temporal_up = temporal_up
|
||||
self.spatial_up = spatial_up
|
||||
self.temporal_ratio = 2 if temporal_up else 1
|
||||
self.spatial_ratio = 2 if spatial_up else 1
|
||||
self.slicing = slicing
|
||||
|
||||
upscale_ratio = (self.spatial_ratio**2) * self.temporal_ratio
|
||||
self.upscale_conv = nn.Conv3d(
|
||||
self.channels, self.channels * upscale_ratio, kernel_size=1, padding=0
|
||||
)
|
||||
identity = (
|
||||
torch.eye(self.channels).repeat(upscale_ratio, 1).reshape_as(self.upscale_conv.weight)
|
||||
)
|
||||
|
||||
self.upscale_conv.weight.data.copy_(identity)
|
||||
nn.init.zeros_(self.upscale_conv.bias)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
memory_state: MemoryState,
|
||||
) -> torch.FloatTensor:
|
||||
return gradient_checkpointing(
|
||||
self.custom_forward,
|
||||
hidden_states,
|
||||
memory_state,
|
||||
enabled=self.training and self.gradient_checkpointing,
|
||||
)
|
||||
|
||||
def custom_forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
memory_state: MemoryState,
|
||||
) -> torch.FloatTensor:
|
||||
assert hidden_states.shape[1] == self.channels
|
||||
|
||||
if self.slicing:
|
||||
split_size = hidden_states.size(2) // 2
|
||||
hidden_states = list(
|
||||
hidden_states.split([split_size, hidden_states.size(2) - split_size], dim=2)
|
||||
)
|
||||
else:
|
||||
hidden_states = [hidden_states]
|
||||
|
||||
for i in range(len(hidden_states)):
|
||||
hidden_states[i] = self.upscale_conv(hidden_states[i])
|
||||
hidden_states[i] = rearrange(
|
||||
hidden_states[i],
|
||||
"b (x y z c) f h w -> b c (f z) (h x) (w y)",
|
||||
x=self.spatial_ratio,
|
||||
y=self.spatial_ratio,
|
||||
z=self.temporal_ratio,
|
||||
)
|
||||
|
||||
# [Overridden] For causal temporal conv
|
||||
if self.temporal_up and memory_state != MemoryState.ACTIVE:
|
||||
hidden_states[0] = remove_head(hidden_states[0])
|
||||
|
||||
if self.slicing:
|
||||
hidden_states = self.conv(hidden_states, memory_state=memory_state)
|
||||
return torch.cat(hidden_states, dim=2)
|
||||
else:
|
||||
return self.conv(hidden_states[0], memory_state=memory_state)
|
||||
|
||||
|
||||
class Downsample3D(nn.Module):
|
||||
"""A 3D downsampling layer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
inflation_mode: _inflation_mode_t = "tail",
|
||||
temporal_down: bool = False,
|
||||
spatial_down: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.temporal_down = temporal_down
|
||||
self.spatial_down = spatial_down
|
||||
|
||||
self.temporal_ratio = 2 if temporal_down else 1
|
||||
self.spatial_ratio = 2 if spatial_down else 1
|
||||
|
||||
self.temporal_kernel = 3 if temporal_down else 1
|
||||
self.spatial_kernel = 3 if spatial_down else 1
|
||||
|
||||
self.conv = init_causal_conv3d(
|
||||
self.channels,
|
||||
self.channels,
|
||||
kernel_size=(self.temporal_kernel, self.spatial_kernel, self.spatial_kernel),
|
||||
stride=(self.temporal_ratio, self.spatial_ratio, self.spatial_ratio),
|
||||
padding=((1 if self.temporal_down else 0), 0, 0),
|
||||
inflation_mode=inflation_mode,
|
||||
)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
memory_state: MemoryState,
|
||||
) -> torch.FloatTensor:
|
||||
return gradient_checkpointing(
|
||||
self.custom_forward,
|
||||
hidden_states,
|
||||
memory_state,
|
||||
enabled=self.training and self.gradient_checkpointing,
|
||||
)
|
||||
|
||||
def custom_forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
memory_state: MemoryState,
|
||||
) -> torch.FloatTensor:
|
||||
|
||||
assert hidden_states.shape[1] == self.channels
|
||||
|
||||
if self.spatial_down:
|
||||
hidden_states = safe_pad_operation(hidden_states, (0, 1, 0, 1), mode="constant", value=0)
|
||||
|
||||
hidden_states = self.conv(hidden_states, memory_state=memory_state)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class ResnetBlock3D(ResnetBlock2D):
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
inflation_mode: _inflation_mode_t = "tail",
|
||||
time_receptive_field: _receptive_field_t = "half",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.conv1 = init_causal_conv3d(
|
||||
self.in_channels,
|
||||
self.out_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
inflation_mode=inflation_mode,
|
||||
)
|
||||
|
||||
self.conv2 = init_causal_conv3d(
|
||||
self.out_channels,
|
||||
self.out_channels,
|
||||
kernel_size=(1, 3, 3) if time_receptive_field == "half" else (3, 3, 3),
|
||||
stride=1,
|
||||
padding=(0, 1, 1) if time_receptive_field == "half" else (1, 1, 1),
|
||||
inflation_mode=inflation_mode,
|
||||
)
|
||||
|
||||
if self.use_in_shortcut:
|
||||
self.conv_shortcut = init_causal_conv3d(
|
||||
self.in_channels,
|
||||
self.out_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
bias=(self.conv_shortcut.bias is not None),
|
||||
inflation_mode=inflation_mode,
|
||||
)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor, memory_state: MemoryState = MemoryState.UNSET):
|
||||
return gradient_checkpointing(
|
||||
self.custom_forward,
|
||||
input_tensor,
|
||||
memory_state,
|
||||
enabled=self.training and self.gradient_checkpointing,
|
||||
)
|
||||
|
||||
def custom_forward(
|
||||
self, input_tensor: torch.Tensor, memory_state: MemoryState = MemoryState.UNSET
|
||||
):
|
||||
assert memory_state != MemoryState.UNSET
|
||||
hidden_states = input_tensor
|
||||
|
||||
hidden_states = causal_norm_wrapper(self.norm1, hidden_states)
|
||||
hidden_states = self.nonlinearity(hidden_states)
|
||||
hidden_states = self.conv1(hidden_states, memory_state=memory_state)
|
||||
|
||||
hidden_states = causal_norm_wrapper(self.norm2, hidden_states)
|
||||
hidden_states = self.nonlinearity(hidden_states)
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
hidden_states = self.conv2(hidden_states, memory_state=memory_state)
|
||||
|
||||
if self.conv_shortcut is not None:
|
||||
input_tensor = self.conv_shortcut(input_tensor, memory_state=memory_state)
|
||||
|
||||
output_tensor = input_tensor + hidden_states
|
||||
|
||||
return output_tensor
|
||||
|
||||
|
||||
class DownEncoderBlock3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
add_downsample: bool = True,
|
||||
inflation_mode: _inflation_mode_t = "tail",
|
||||
time_receptive_field: _receptive_field_t = "half",
|
||||
temporal_down: bool = True,
|
||||
spatial_down: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
resnets = []
|
||||
|
||||
for i in range(num_layers):
|
||||
in_channels = in_channels if i == 0 else out_channels
|
||||
resnets.append(
|
||||
ResnetBlock3D(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
dropout=dropout,
|
||||
inflation_mode=inflation_mode,
|
||||
time_receptive_field=time_receptive_field,
|
||||
)
|
||||
)
|
||||
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
|
||||
self.downsamplers = None
|
||||
if add_downsample:
|
||||
# Todo: Refactor this line before V5 Image VAE Training.
|
||||
self.downsamplers = nn.ModuleList(
|
||||
[
|
||||
Downsample3D(
|
||||
channels=out_channels,
|
||||
inflation_mode=inflation_mode,
|
||||
temporal_down=temporal_down,
|
||||
spatial_down=spatial_down,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, hidden_states: torch.FloatTensor, memory_state: MemoryState
|
||||
) -> torch.FloatTensor:
|
||||
for resnet in self.resnets:
|
||||
hidden_states = resnet(hidden_states, memory_state=memory_state)
|
||||
|
||||
if self.downsamplers is not None:
|
||||
for downsampler in self.downsamplers:
|
||||
hidden_states = downsampler(hidden_states, memory_state=memory_state)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class UpDecoderBlock3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
add_upsample: bool = True,
|
||||
inflation_mode: _inflation_mode_t = "tail",
|
||||
time_receptive_field: _receptive_field_t = "half",
|
||||
temporal_up: bool = True,
|
||||
spatial_up: bool = True,
|
||||
slicing: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
resnets = []
|
||||
|
||||
for i in range(num_layers):
|
||||
input_channels = in_channels if i == 0 else out_channels
|
||||
|
||||
resnets.append(
|
||||
ResnetBlock3D(
|
||||
in_channels=input_channels,
|
||||
out_channels=out_channels,
|
||||
dropout=dropout,
|
||||
inflation_mode=inflation_mode,
|
||||
time_receptive_field=time_receptive_field,
|
||||
)
|
||||
)
|
||||
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
|
||||
self.upsamplers = None
|
||||
# Todo: Refactor this line before V5 Image VAE Training.
|
||||
if add_upsample:
|
||||
self.upsamplers = nn.ModuleList(
|
||||
[
|
||||
Upsample3D(
|
||||
channels=out_channels,
|
||||
inflation_mode=inflation_mode,
|
||||
temporal_up=temporal_up,
|
||||
spatial_up=spatial_up,
|
||||
slicing=slicing,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, hidden_states: torch.FloatTensor, memory_state: MemoryState
|
||||
) -> torch.FloatTensor:
|
||||
for resnet in self.resnets:
|
||||
hidden_states = resnet(hidden_states, memory_state=memory_state)
|
||||
|
||||
if self.upsamplers is not None:
|
||||
for upsampler in self.upsamplers:
|
||||
hidden_states = upsampler(hidden_states, memory_state=memory_state)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class UNetMidBlock3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
dropout: float = 0.0,
|
||||
inflation_mode: _inflation_mode_t = "tail",
|
||||
time_receptive_field: _receptive_field_t = "half",
|
||||
):
|
||||
super().__init__()
|
||||
self.resnets = nn.ModuleList(
|
||||
[
|
||||
ResnetBlock3D(
|
||||
in_channels=channels,
|
||||
out_channels=channels,
|
||||
dropout=dropout,
|
||||
inflation_mode=inflation_mode,
|
||||
time_receptive_field=time_receptive_field,
|
||||
),
|
||||
ResnetBlock3D(
|
||||
in_channels=channels,
|
||||
out_channels=channels,
|
||||
dropout=dropout,
|
||||
inflation_mode=inflation_mode,
|
||||
time_receptive_field=time_receptive_field,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor, memory_state: MemoryState):
|
||||
for resnet in self.resnets:
|
||||
hidden_states = resnet(hidden_states, memory_state)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class Encoder3D(nn.Module):
|
||||
r"""
|
||||
The `Encoder` layer of a variational autoencoder that encodes
|
||||
its input into a latent representation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int = 3,
|
||||
out_channels: int = 3,
|
||||
block_out_channels: Tuple[int, ...] = (64,),
|
||||
layers_per_block: int = 2,
|
||||
double_z: bool = True,
|
||||
temporal_down_num: int = 2,
|
||||
inflation_mode: _inflation_mode_t = "tail",
|
||||
time_receptive_field: _receptive_field_t = "half",
|
||||
selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
|
||||
):
|
||||
super().__init__()
|
||||
self.layers_per_block = layers_per_block
|
||||
|
||||
self.temporal_down_num = temporal_down_num
|
||||
|
||||
self.conv_in = init_causal_conv3d(
|
||||
in_channels,
|
||||
block_out_channels[0],
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
inflation_mode=inflation_mode,
|
||||
)
|
||||
|
||||
self.down_blocks = nn.ModuleList([])
|
||||
|
||||
# down
|
||||
output_channel = block_out_channels[0]
|
||||
for i in range(len(block_out_channels)):
|
||||
input_channel = output_channel
|
||||
output_channel = block_out_channels[i]
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
is_temporal_down_block = i >= len(block_out_channels) - self.temporal_down_num - 1
|
||||
# Note: take the last one
|
||||
|
||||
down_block = DownEncoderBlock3D(
|
||||
num_layers=self.layers_per_block,
|
||||
in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
add_downsample=not is_final_block,
|
||||
temporal_down=is_temporal_down_block,
|
||||
spatial_down=True,
|
||||
inflation_mode=inflation_mode,
|
||||
time_receptive_field=time_receptive_field,
|
||||
)
|
||||
self.down_blocks.append(down_block)
|
||||
|
||||
# mid
|
||||
self.mid_block = UNetMidBlock3D(
|
||||
channels=block_out_channels[-1],
|
||||
inflation_mode=inflation_mode,
|
||||
time_receptive_field=time_receptive_field,
|
||||
)
|
||||
|
||||
# out
|
||||
self.conv_norm_out = nn.GroupNorm(
|
||||
num_channels=block_out_channels[-1], num_groups=32, eps=1e-6
|
||||
)
|
||||
self.conv_act = nn.SiLU()
|
||||
|
||||
conv_out_channels = 2 * out_channels if double_z else out_channels
|
||||
self.conv_out = init_causal_conv3d(
|
||||
block_out_channels[-1], conv_out_channels, 3, padding=1, inflation_mode=inflation_mode
|
||||
)
|
||||
|
||||
assert len(selective_checkpointing) == len(self.down_blocks)
|
||||
self.set_gradient_checkpointing(selective_checkpointing)
|
||||
|
||||
def set_gradient_checkpointing(self, checkpointing_types):
|
||||
gradient_checkpointing = []
|
||||
for down_block, sac_type in zip(self.down_blocks, checkpointing_types):
|
||||
if sac_type == "coarse":
|
||||
gradient_checkpointing.append(True)
|
||||
elif sac_type == "fine":
|
||||
for n, m in down_block.named_modules():
|
||||
if hasattr(m, "gradient_checkpointing"):
|
||||
m.gradient_checkpointing = True
|
||||
logger.debug(f"set gradient_checkpointing: {n}")
|
||||
gradient_checkpointing.append(False)
|
||||
else:
|
||||
gradient_checkpointing.append(False)
|
||||
self.gradient_checkpointing = gradient_checkpointing
|
||||
logger.info(f"[Encoder3D] gradient_checkpointing: {checkpointing_types}")
|
||||
|
||||
def forward(self, sample: torch.FloatTensor, memory_state: MemoryState) -> torch.FloatTensor:
|
||||
r"""The forward method of the `Encoder` class."""
|
||||
sample = self.conv_in(sample, memory_state=memory_state)
|
||||
# down
|
||||
for down_block, sac in zip(self.down_blocks, self.gradient_checkpointing):
|
||||
sample = gradient_checkpointing(
|
||||
down_block,
|
||||
sample,
|
||||
memory_state=memory_state,
|
||||
enabled=self.training and sac,
|
||||
)
|
||||
|
||||
# middle
|
||||
sample = self.mid_block(sample, memory_state=memory_state)
|
||||
|
||||
# post-process
|
||||
sample = causal_norm_wrapper(self.conv_norm_out, sample)
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample, memory_state=memory_state)
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class Decoder3D(nn.Module):
|
||||
r"""
|
||||
The `Decoder` layer of a variational autoencoder that
|
||||
decodes its latent representation into an output sample.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int = 3,
|
||||
out_channels: int = 3,
|
||||
block_out_channels: Tuple[int, ...] = (64,),
|
||||
layers_per_block: int = 2,
|
||||
inflation_mode: _inflation_mode_t = "tail",
|
||||
time_receptive_field: _receptive_field_t = "half",
|
||||
temporal_up_num: int = 2,
|
||||
slicing_up_num: int = 0,
|
||||
selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
|
||||
):
|
||||
super().__init__()
|
||||
self.layers_per_block = layers_per_block
|
||||
self.temporal_up_num = temporal_up_num
|
||||
|
||||
self.conv_in = init_causal_conv3d(
|
||||
in_channels,
|
||||
block_out_channels[-1],
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
inflation_mode=inflation_mode,
|
||||
)
|
||||
|
||||
self.up_blocks = nn.ModuleList([])
|
||||
|
||||
# mid
|
||||
self.mid_block = UNetMidBlock3D(
|
||||
channels=block_out_channels[-1],
|
||||
inflation_mode=inflation_mode,
|
||||
time_receptive_field=time_receptive_field,
|
||||
)
|
||||
|
||||
# up
|
||||
reversed_block_out_channels = list(reversed(block_out_channels))
|
||||
output_channel = reversed_block_out_channels[0]
|
||||
for i in range(len(reversed_block_out_channels)):
|
||||
prev_output_channel = output_channel
|
||||
output_channel = reversed_block_out_channels[i]
|
||||
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
is_temporal_up_block = i < self.temporal_up_num
|
||||
is_slicing_up_block = i >= len(block_out_channels) - slicing_up_num
|
||||
# Note: Keep symmetric
|
||||
|
||||
up_block = UpDecoderBlock3D(
|
||||
num_layers=self.layers_per_block + 1,
|
||||
in_channels=prev_output_channel,
|
||||
out_channels=output_channel,
|
||||
add_upsample=not is_final_block,
|
||||
temporal_up=is_temporal_up_block,
|
||||
slicing=is_slicing_up_block,
|
||||
inflation_mode=inflation_mode,
|
||||
time_receptive_field=time_receptive_field,
|
||||
)
|
||||
self.up_blocks.append(up_block)
|
||||
|
||||
# out
|
||||
self.conv_norm_out = nn.GroupNorm(
|
||||
num_channels=block_out_channels[0], num_groups=32, eps=1e-6
|
||||
)
|
||||
self.conv_act = nn.SiLU()
|
||||
self.conv_out = init_causal_conv3d(
|
||||
block_out_channels[0], out_channels, 3, padding=1, inflation_mode=inflation_mode
|
||||
)
|
||||
|
||||
assert len(selective_checkpointing) == len(self.up_blocks)
|
||||
self.set_gradient_checkpointing(selective_checkpointing)
|
||||
|
||||
def set_gradient_checkpointing(self, checkpointing_types):
|
||||
gradient_checkpointing = []
|
||||
for up_block, sac_type in zip(self.up_blocks, checkpointing_types):
|
||||
if sac_type == "coarse":
|
||||
gradient_checkpointing.append(True)
|
||||
elif sac_type == "fine":
|
||||
for n, m in up_block.named_modules():
|
||||
if hasattr(m, "gradient_checkpointing"):
|
||||
m.gradient_checkpointing = True
|
||||
logger.debug(f"set gradient_checkpointing: {n}")
|
||||
gradient_checkpointing.append(False)
|
||||
else:
|
||||
gradient_checkpointing.append(False)
|
||||
self.gradient_checkpointing = gradient_checkpointing
|
||||
logger.info(f"[Decoder3D] gradient_checkpointing: {checkpointing_types}")
|
||||
|
||||
def forward(self, sample: torch.FloatTensor, memory_state: MemoryState) -> torch.FloatTensor:
|
||||
r"""The forward method of the `Decoder` class."""
|
||||
|
||||
sample = self.conv_in(sample, memory_state=memory_state)
|
||||
|
||||
# middle
|
||||
sample = self.mid_block(sample, memory_state=memory_state)
|
||||
|
||||
# up
|
||||
for up_block, sac in zip(self.up_blocks, self.gradient_checkpointing):
|
||||
sample = gradient_checkpointing(
|
||||
up_block,
|
||||
sample,
|
||||
memory_state=memory_state,
|
||||
enabled=self.training and sac,
|
||||
)
|
||||
|
||||
# post-process
|
||||
sample = causal_norm_wrapper(self.conv_norm_out, sample)
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample, memory_state=memory_state)
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class VideoAutoencoderKL(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int = 3,
|
||||
out_channels: int = 3,
|
||||
block_out_channels: Tuple[int] = (64,),
|
||||
layers_per_block: int = 1,
|
||||
latent_channels: int = 4,
|
||||
use_quant_conv: bool = True,
|
||||
use_post_quant_conv: bool = True,
|
||||
enc_selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
|
||||
dec_selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
|
||||
temporal_scale_num: int = 0,
|
||||
slicing_up_num: int = 0,
|
||||
inflation_mode: _inflation_mode_t = "tail",
|
||||
time_receptive_field: _receptive_field_t = "half",
|
||||
slicing_sample_min_size: int = None,
|
||||
spatial_downsample_factor: int = 16,
|
||||
temporal_downsample_factor: int = 8,
|
||||
freeze_encoder: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.spatial_downsample_factor = spatial_downsample_factor
|
||||
self.temporal_downsample_factor = temporal_downsample_factor
|
||||
self.freeze_encoder = freeze_encoder
|
||||
if slicing_sample_min_size is None:
|
||||
slicing_sample_min_size = temporal_downsample_factor
|
||||
self.slicing_sample_min_size = slicing_sample_min_size
|
||||
self.slicing_latent_min_size = slicing_sample_min_size // (2**temporal_scale_num)
|
||||
|
||||
# pass init params to Encoder
|
||||
self.encoder = Encoder3D(
|
||||
in_channels=in_channels,
|
||||
out_channels=latent_channels,
|
||||
block_out_channels=block_out_channels,
|
||||
layers_per_block=layers_per_block,
|
||||
double_z=True,
|
||||
temporal_down_num=temporal_scale_num,
|
||||
selective_checkpointing=enc_selective_checkpointing,
|
||||
inflation_mode=inflation_mode,
|
||||
time_receptive_field=time_receptive_field,
|
||||
)
|
||||
|
||||
# pass init params to Decoder
|
||||
self.decoder = Decoder3D(
|
||||
in_channels=latent_channels,
|
||||
out_channels=out_channels,
|
||||
block_out_channels=block_out_channels,
|
||||
layers_per_block=layers_per_block,
|
||||
# [Override] add temporal_up_num parameter
|
||||
temporal_up_num=temporal_scale_num,
|
||||
slicing_up_num=slicing_up_num,
|
||||
selective_checkpointing=dec_selective_checkpointing,
|
||||
inflation_mode=inflation_mode,
|
||||
time_receptive_field=time_receptive_field,
|
||||
)
|
||||
|
||||
self.quant_conv = (
|
||||
init_causal_conv3d(
|
||||
in_channels=2 * latent_channels,
|
||||
out_channels=2 * latent_channels,
|
||||
kernel_size=1,
|
||||
inflation_mode=inflation_mode,
|
||||
)
|
||||
if use_quant_conv
|
||||
else None
|
||||
)
|
||||
self.post_quant_conv = (
|
||||
init_causal_conv3d(
|
||||
in_channels=latent_channels,
|
||||
out_channels=latent_channels,
|
||||
kernel_size=1,
|
||||
inflation_mode=inflation_mode,
|
||||
)
|
||||
if use_post_quant_conv
|
||||
else None
|
||||
)
|
||||
|
||||
self.use_slicing = False
|
||||
|
||||
def enable_slicing(self):
|
||||
self.use_slicing = True
|
||||
|
||||
def disable_slicing(self):
|
||||
self.use_slicing = False
|
||||
|
||||
def encode(self, x: torch.FloatTensor) -> CausalEncoderOutput:
|
||||
if x.ndim == 4:
|
||||
x = x.unsqueeze(2)
|
||||
h = self.slicing_encode(x)
|
||||
p = DiagonalGaussianDistribution(h)
|
||||
z = p.sample()
|
||||
return CausalEncoderOutput(z, p)
|
||||
|
||||
def decode(self, z: torch.FloatTensor) -> CausalDecoderOutput:
|
||||
if z.ndim == 4:
|
||||
z = z.unsqueeze(2)
|
||||
x = self.slicing_decode(z)
|
||||
return CausalDecoderOutput(x)
|
||||
|
||||
def _encode(self, x: torch.Tensor, memory_state: MemoryState) -> torch.Tensor:
|
||||
x = causal_conv_slice_inputs(x, self.slicing_sample_min_size, memory_state=memory_state)
|
||||
h = self.encoder(x, memory_state=memory_state)
|
||||
h = self.quant_conv(h, memory_state=memory_state) if self.quant_conv is not None else h
|
||||
h = causal_conv_gather_outputs(h)
|
||||
return h
|
||||
|
||||
def _decode(self, z: torch.Tensor, memory_state: MemoryState) -> torch.Tensor:
|
||||
z = causal_conv_slice_inputs(z, self.slicing_latent_min_size, memory_state=memory_state)
|
||||
z = (
|
||||
self.post_quant_conv(z, memory_state=memory_state)
|
||||
if self.post_quant_conv is not None
|
||||
else z
|
||||
)
|
||||
x = self.decoder(z, memory_state=memory_state)
|
||||
x = causal_conv_gather_outputs(x)
|
||||
return x
|
||||
|
||||
def slicing_encode(self, x: torch.Tensor) -> torch.Tensor:
|
||||
sp_size = 1
|
||||
if self.use_slicing and (x.shape[2] - 1) > self.slicing_sample_min_size * sp_size:
|
||||
x_slices = x[:, :, 1:].split(split_size=self.slicing_sample_min_size * sp_size, dim=2)
|
||||
encoded_slices = [
|
||||
self._encode(
|
||||
torch.cat((x[:, :, :1], x_slices[0]), dim=2),
|
||||
memory_state=MemoryState.INITIALIZING,
|
||||
)
|
||||
]
|
||||
for x_idx in range(1, len(x_slices)):
|
||||
encoded_slices.append(
|
||||
self._encode(x_slices[x_idx], memory_state=MemoryState.ACTIVE)
|
||||
)
|
||||
return torch.cat(encoded_slices, dim=2)
|
||||
else:
|
||||
return self._encode(x, memory_state=MemoryState.DISABLED)
|
||||
|
||||
def slicing_decode(self, z: torch.Tensor) -> torch.Tensor:
|
||||
sp_size = 1
|
||||
if self.use_slicing and (z.shape[2] - 1) > self.slicing_latent_min_size * sp_size:
|
||||
z_slices = z[:, :, 1:].split(split_size=self.slicing_latent_min_size * sp_size, dim=2)
|
||||
decoded_slices = [
|
||||
self._decode(
|
||||
torch.cat((z[:, :, :1], z_slices[0]), dim=2),
|
||||
memory_state=MemoryState.INITIALIZING,
|
||||
)
|
||||
]
|
||||
for z_idx in range(1, len(z_slices)):
|
||||
decoded_slices.append(
|
||||
self._decode(z_slices[z_idx], memory_state=MemoryState.ACTIVE)
|
||||
)
|
||||
return torch.cat(decoded_slices, dim=2)
|
||||
else:
|
||||
return self._decode(z, memory_state=MemoryState.DISABLED)
|
||||
|
||||
def forward(self, x: torch.FloatTensor) -> CausalAutoencoderOutput:
|
||||
with torch.no_grad() if self.freeze_encoder else nullcontext():
|
||||
z, p = self.encode(x)
|
||||
x = self.decode(z).sample
|
||||
return CausalAutoencoderOutput(x, z, p)
|
||||
|
||||
def preprocess(self, x: torch.Tensor):
|
||||
# x should in [B, C, T, H, W], [B, C, H, W]
|
||||
assert x.ndim == 4 or x.size(2) % self.temporal_downsample_factor == 1
|
||||
return x
|
||||
|
||||
def postprocess(self, x: torch.Tensor):
|
||||
# x should in [B, C, T, H, W], [B, C, H, W]
|
||||
return x
|
||||
|
||||
def set_causal_slicing(
|
||||
self,
|
||||
*,
|
||||
split_size: Optional[int],
|
||||
memory_device: _memory_device_t,
|
||||
):
|
||||
assert (
|
||||
split_size is None or memory_device is not None
|
||||
), "if split_size is set, memory_device must not be None."
|
||||
if split_size is not None:
|
||||
self.enable_slicing()
|
||||
self.slicing_sample_min_size = split_size
|
||||
self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
|
||||
else:
|
||||
self.disable_slicing()
|
||||
for module in self.modules():
|
||||
if isinstance(module, InflatedCausalConv3d):
|
||||
module.set_memory_device(memory_device)
|
||||
|
||||
def set_memory_limit(self, conv_max_mem: Optional[float], norm_max_mem: Optional[float]):
|
||||
set_norm_limit(norm_max_mem)
|
||||
for m in self.modules():
|
||||
if isinstance(m, InflatedCausalConv3d):
|
||||
m.set_memory_limit(conv_max_mem if conv_max_mem is not None else float("inf"))
|
||||
|
||||
|
||||
class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
|
||||
def __init__(
|
||||
self, *args, spatial_downsample_factor: int, temporal_downsample_factor: int, **kwargs
|
||||
):
|
||||
self.spatial_downsample_factor = spatial_downsample_factor
|
||||
self.temporal_downsample_factor = temporal_downsample_factor
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def forward(self, x) -> CausalAutoencoderOutput:
|
||||
z, _, p = self.encode(x)
|
||||
x, _ = self.decode(z)
|
||||
return CausalAutoencoderOutput(x, z, None, p)
|
||||
|
||||
def encode(self, x) -> CausalEncoderOutput:
|
||||
if x.ndim == 4:
|
||||
x = x.unsqueeze(2)
|
||||
p = super().encode(x).latent_dist
|
||||
z = p.sample().squeeze(2)
|
||||
return CausalEncoderOutput(z, None, p)
|
||||
|
||||
def decode(self, z) -> CausalDecoderOutput:
|
||||
if z.ndim == 4:
|
||||
z = z.unsqueeze(2)
|
||||
x = super().decode(z).sample.squeeze(2)
|
||||
return CausalDecoderOutput(x, None)
|
||||
|
||||
def preprocess(self, x):
|
||||
# x should in [B, C, T, H, W], [B, C, H, W]
|
||||
assert x.ndim == 4 or x.size(2) % 4 == 1
|
||||
return x
|
||||
|
||||
def postprocess(self, x):
|
||||
# x should in [B, C, T, H, W], [B, C, H, W]
|
||||
return x
|
||||
|
||||
def set_causal_slicing(
|
||||
self,
|
||||
*,
|
||||
split_size: Optional[int],
|
||||
memory_device: Optional[Literal["cpu", "same"]],
|
||||
):
|
||||
assert (
|
||||
split_size is None or memory_device is not None
|
||||
), "if split_size is set, memory_device must not be None."
|
||||
if split_size is not None:
|
||||
self.enable_slicing()
|
||||
else:
|
||||
self.disable_slicing()
|
||||
self.slicing_sample_min_size = split_size
|
||||
if split_size is not None:
|
||||
self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
|
||||
for module in self.modules():
|
||||
if isinstance(module, InflatedCausalConv3d):
|
||||
module.set_memory_device(memory_device)
|
||||
@@ -36,6 +36,8 @@ def upscale_image(model_name:str, image_path:str):
|
||||
runner=runner,
|
||||
images=image_tensor,
|
||||
cfg_scale=cfg,
|
||||
cfg_rescale=0.0,
|
||||
steps=1,
|
||||
seed=seed,
|
||||
res_w=resolution,
|
||||
batch_size=1,
|
||||
|
||||
@@ -631,9 +631,6 @@ def create_settings(cmd_opts):
|
||||
"detailer_unload": OptionInfo(False, "Move detailer model to CPU when complete"),
|
||||
"detailer_augment": OptionInfo(False, "Detailer use model augment"),
|
||||
|
||||
"postprocessing_sep_seedvr": OptionInfo("<h2>SeedVR</h2>", "", gr.HTML),
|
||||
"seedvr_cfg_scale": OptionInfo(3.5, "SeedVR CFG Scale", gr.Slider, {"minimum": 1, "maximum": 15, "step": 1}),
|
||||
|
||||
"postprocessing_sep_upscalers": OptionInfo("<h2>Upscaling</h2>", "", gr.HTML),
|
||||
"upscaler_unload": OptionInfo(False, "Unload upscaler after processing"),
|
||||
"upscaler_latent_steps": OptionInfo(20, "Upscaler latent steps", gr.Slider, {"minimum": 4, "maximum": 100, "step": 1}),
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import gradio as gr
|
||||
from modules import scripts_postprocessing, devices
|
||||
|
||||
|
||||
class ScriptPixelArt(scripts_postprocessing.ScriptPostprocessing):
|
||||
name = "PixelArt"
|
||||
order = 30000
|
||||
@@ -13,6 +14,7 @@ class ScriptPixelArt(scripts_postprocessing.ScriptPostprocessing):
|
||||
with gr.Row():
|
||||
pixelart_block_size = gr.Slider(minimum=2, maximum=64, step=1, value=8, label="PixelArt block size", elem_id="extras_pixelart_block_size")
|
||||
pixelart_edge_block_size = gr.Slider(minimum=2, maximum=64, step=1, value=4, label="Edge block size", elem_id="extras_pixelart_edge_block_size")
|
||||
with gr.Row():
|
||||
pixelart_image_weight = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=1.0, label="Edge image weight", elem_id="extras_pixelart_image_weight")
|
||||
pixelart_sharpen_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.1, label="PixelArt sharpen", elem_id="extras_pixelart_sharpen_amount")
|
||||
return {
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
import gradio as gr
|
||||
from modules import scripts_postprocessing
|
||||
|
||||
|
||||
class ScriptSeedVR(scripts_postprocessing.ScriptPostprocessing):
|
||||
name = "SeedVR"
|
||||
|
||||
def ui(self):
|
||||
from modules.postprocess.seedvr_model import MODELS_MAP
|
||||
with gr.Accordion(self.name, open = False, elem_id="postprocess_seedvr_accordion"):
|
||||
with gr.Row():
|
||||
seedvr_enabled = gr.Checkbox(label="Enable SeedVR", value=False, elem_id="extras_seedvr_enabled")
|
||||
with gr.Row():
|
||||
seedvr_selected = gr.Dropdown(label="SeedVR model", choices=list(MODELS_MAP.keys()), value=list(MODELS_MAP.keys())[0], elem_id="extras_seedvr_model")
|
||||
with gr.Row():
|
||||
seedvr_scale = gr.Slider(minimum=1, maximum=16, step=0.1, value=2, label="SeedVR scale", elem_id="extras_seedvr_scale")
|
||||
seedvr_seed = gr.Number(step=1, value=-1, label="SeedVR seed", elem_id="extras_seedvr_seed")
|
||||
seedvr_steps = gr.Number(step=1, value=1, minimum=1, maximum=99, label="SeedVR steps", elem_id="extras_seedvr_steps", visible=False)
|
||||
with gr.Row():
|
||||
seedvr_cfg_scale = gr.Slider(minimum=0.0, maximum=15.0, step=0.01, value=3.5, label="SeedVR guidance scale", elem_id="extras_seedvr_cfg_scale")
|
||||
seedvr_cfg_rescale = gr.Slider(minimum=0.0, maximum=15.0, step=0.01, value=0.0, label="SeedVR guidance rescale", elem_id="extras_seedvr_cfg_rescale")
|
||||
with gr.Row():
|
||||
seedvr_tile_size = gr.Slider(minimum=64, maximum=4096, step=8, value=1024, label="SeedVR tile size", elem_id="extras_seedvr_tile_size")
|
||||
seedvr_tile_overlap = gr.Slider(minimum=0, maximum=1.0, step=0.01, value=0.25, label="SeedVR tile overlap", elem_id="extras_seedvr_tile_overlap")
|
||||
return {
|
||||
"seedvr_enabled": seedvr_enabled,
|
||||
"seedvr_selected": seedvr_selected,
|
||||
"seedvr_scale": seedvr_scale,
|
||||
"seedvr_seed": seedvr_seed,
|
||||
"seedvr_steps": seedvr_steps,
|
||||
"seedvr_cfg_scale": seedvr_cfg_scale,
|
||||
"seedvr_cfg_rescale": seedvr_cfg_rescale,
|
||||
"seedvr_tile_size": seedvr_tile_size,
|
||||
"seedvr_tile_overlap": seedvr_tile_overlap,
|
||||
}
|
||||
|
||||
def process(self,
|
||||
pp: scripts_postprocessing.PostprocessedImage,
|
||||
seedvr_enabled: bool,
|
||||
seedvr_selected: str,
|
||||
seedvr_scale: int,
|
||||
seedvr_seed: int,
|
||||
seedvr_steps: int,
|
||||
seedvr_cfg_scale: float,
|
||||
seedvr_cfg_rescale: float,
|
||||
seedvr_tile_size: int,
|
||||
seedvr_tile_overlap: float
|
||||
): # pylint: disable=arguments-differ
|
||||
if not seedvr_enabled:
|
||||
return
|
||||
from modules import shared, upscaler
|
||||
from modules.logger import log
|
||||
image = pp.image
|
||||
instance: upscaler.UpscalerData = next(iter([x for x in shared.sd_upscalers if x.name == seedvr_selected]), None)
|
||||
scaler: upscaler.Upscaler = instance.scaler
|
||||
|
||||
log.info(f'Upscaler: type="SeedVR" model="{seedvr_selected}" scale={seedvr_scale} seed={seedvr_seed} steps={seedvr_steps} cfg_scale={seedvr_cfg_scale} cfg_rescale={seedvr_cfg_rescale} tile_size={seedvr_tile_size} tile_overlap={seedvr_tile_overlap}')
|
||||
|
||||
jobid = shared.state.begin('Upscale')
|
||||
|
||||
scaler.scale = float(seedvr_scale)
|
||||
upscaled = scaler.do_upscale(image,
|
||||
seedvr_selected,
|
||||
cfg_scale=seedvr_cfg_scale,
|
||||
cfg_rescale=seedvr_cfg_rescale,
|
||||
steps=seedvr_steps,
|
||||
seed=seedvr_seed,
|
||||
tile_size=seedvr_tile_size,
|
||||
tile_overlap=seedvr_tile_overlap,
|
||||
)
|
||||
shared.state.end(jobid)
|
||||
|
||||
pp.image = upscaled
|
||||
pp.info["SeedVR"] = f"Scale={seedvr_scale} Seed={seedvr_seed} CFG Scale={seedvr_cfg_scale} CFG Rescale={seedvr_cfg_rescale}"
|
||||
Reference in New Issue
Block a user