explicit seedvr implementation

Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2026-07-18 10:35:07 +02:00
parent 5c47e557c4
commit 42cb4cf489
16 changed files with 154 additions and 972 deletions
+5 -2
View File
@@ -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
+18 -6
View File
@@ -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):
+3 -3
View File
@@ -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
+1 -4
View File
@@ -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)
+24 -5
View File
@@ -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)
+2
View File
@@ -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,
-3
View File
@@ -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}),
+2
View File
@@ -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 {
+74
View File
@@ -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}"