mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
major refactoring of modules
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
@@ -0,0 +1,310 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from collections import OrderedDict
|
||||
from diffusers.models.embeddings import (
|
||||
TimestepEmbedding,
|
||||
Timesteps,
|
||||
)
|
||||
|
||||
|
||||
def conv_nd(dims, *args, **kwargs):
|
||||
"""
|
||||
Create a 1D, 2D, or 3D convolution module.
|
||||
"""
|
||||
if dims == 1:
|
||||
return nn.Conv1d(*args, **kwargs)
|
||||
elif dims == 2:
|
||||
return nn.Conv2d(*args, **kwargs)
|
||||
elif dims == 3:
|
||||
return nn.Conv3d(*args, **kwargs)
|
||||
raise ValueError(f"unsupported dimensions: {dims}")
|
||||
|
||||
|
||||
def avg_pool_nd(dims, *args, **kwargs):
|
||||
"""
|
||||
Create a 1D, 2D, or 3D average pooling module.
|
||||
"""
|
||||
if dims == 1:
|
||||
return nn.AvgPool1d(*args, **kwargs)
|
||||
elif dims == 2:
|
||||
return nn.AvgPool2d(*args, **kwargs)
|
||||
elif dims == 3:
|
||||
return nn.AvgPool3d(*args, **kwargs)
|
||||
raise ValueError(f"unsupported dimensions: {dims}")
|
||||
|
||||
|
||||
def get_parameter_dtype(parameter: torch.nn.Module):
|
||||
try:
|
||||
params = tuple(parameter.parameters())
|
||||
if len(params) > 0:
|
||||
return params[0].dtype
|
||||
|
||||
buffers = tuple(parameter.buffers())
|
||||
if len(buffers) > 0:
|
||||
return buffers[0].dtype
|
||||
|
||||
except StopIteration:
|
||||
# For torch.nn.DataParallel compatibility in PyTorch 1.5
|
||||
|
||||
def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
|
||||
tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
|
||||
return tuples
|
||||
|
||||
gen = parameter._named_members(get_members_fn=find_tensor_attributes)
|
||||
first_tuple = next(gen)
|
||||
return first_tuple[1].dtype
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
"""
|
||||
A downsampling layer with an optional convolution.
|
||||
:param channels: channels in the inputs and outputs.
|
||||
:param use_conv: a bool determining if a convolution is applied.
|
||||
:param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
|
||||
downsampling occurs in the inner-two dimensions.
|
||||
"""
|
||||
|
||||
def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels or channels
|
||||
self.use_conv = use_conv
|
||||
self.dims = dims
|
||||
stride = 2 if dims != 3 else (1, 2, 2)
|
||||
if use_conv:
|
||||
self.op = conv_nd(dims, self.channels, self.out_channels, 3, stride=stride, padding=padding)
|
||||
else:
|
||||
assert self.channels == self.out_channels
|
||||
from torch.nn import MaxUnpool2d
|
||||
self.op = MaxUnpool2d(dims, kernel_size=stride, stride=stride)
|
||||
|
||||
def forward(self, x):
|
||||
assert x.shape[1] == self.channels
|
||||
return self.op(x)
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels or channels
|
||||
self.use_conv = use_conv
|
||||
self.dims = dims
|
||||
stride = 2 if dims != 3 else (1, 2, 2)
|
||||
if use_conv:
|
||||
self.op = nn.ConvTranspose2d(self.channels, self.out_channels, 3, stride=stride, padding=1)
|
||||
else:
|
||||
assert self.channels == self.out_channels
|
||||
self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
|
||||
|
||||
def forward(self, x, output_size):
|
||||
assert x.shape[1] == self.channels
|
||||
return self.op(x, output_size)
|
||||
|
||||
|
||||
class Linear(nn.Module):
|
||||
def __init__(self, temb_channels, out_channels):
|
||||
super(Linear, self).__init__()
|
||||
self.linear = nn.Linear(temb_channels, out_channels)
|
||||
|
||||
def forward(self, x):
|
||||
return self.linear(x)
|
||||
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
|
||||
def __init__(self, in_c, out_c, down, up, ksize=3, sk=False, use_conv=True, enable_timestep=False, temb_channels=None, use_norm=False):
|
||||
super().__init__()
|
||||
self.use_norm = use_norm
|
||||
self.enable_timestep = enable_timestep
|
||||
ps = ksize // 2
|
||||
if in_c != out_c or sk == False:
|
||||
self.in_conv = nn.Conv2d(in_c, out_c, ksize, 1, ps)
|
||||
else:
|
||||
self.in_conv = None
|
||||
self.block1 = nn.Conv2d(out_c, out_c, 3, 1, 1)
|
||||
self.act = nn.ReLU()
|
||||
if use_norm:
|
||||
self.norm1 = nn.GroupNorm(num_groups=32, num_channels=out_c, eps=1e-6, affine=True)
|
||||
self.block2 = nn.Conv2d(out_c, out_c, ksize, 1, ps)
|
||||
if sk == False:
|
||||
self.skep = nn.Conv2d(in_c, out_c, ksize, 1, ps)
|
||||
else:
|
||||
self.skep = None
|
||||
|
||||
self.down = down
|
||||
self.up = up
|
||||
if self.down:
|
||||
self.down_opt = Downsample(in_c, use_conv=use_conv)
|
||||
if self.up:
|
||||
self.up_opt = Upsample(in_c, use_conv=use_conv)
|
||||
if enable_timestep:
|
||||
self.timestep_proj = Linear(temb_channels, out_c)
|
||||
|
||||
|
||||
def forward(self, x, output_size=None, temb=None):
|
||||
if self.down == True:
|
||||
x = self.down_opt(x)
|
||||
if self.up == True:
|
||||
x = self.up_opt(x, output_size)
|
||||
if self.in_conv is not None: # edit
|
||||
x = self.in_conv(x)
|
||||
|
||||
h = self.block1(x)
|
||||
if temb is not None:
|
||||
temb = self.timestep_proj(temb)[:, :, None, None]
|
||||
h = h + temb
|
||||
if self.use_norm:
|
||||
h = self.norm1(h)
|
||||
h = self.act(h)
|
||||
h = self.block2(h)
|
||||
if self.skep is not None:
|
||||
return h + self.skep(x)
|
||||
else:
|
||||
return h + x
|
||||
|
||||
|
||||
class Adapter_XL(nn.Module):
|
||||
|
||||
def __init__(self, in_channels=[1280, 640, 320], out_channels=[1280, 1280, 640], nums_rb=3, ksize=3, sk=True, use_conv=False, use_zero_conv=True,
|
||||
enable_timestep=False, use_norm=False, temb_channels=None, fusion_type='ADD'):
|
||||
super(Adapter_XL, self).__init__()
|
||||
self.channels = in_channels
|
||||
self.nums_rb = nums_rb
|
||||
self.body = []
|
||||
self.out = []
|
||||
self.use_zero_conv = use_zero_conv
|
||||
self.fusion_type = fusion_type
|
||||
self.gamma = []
|
||||
self.beta = []
|
||||
self.norm = []
|
||||
if fusion_type == "SPADE":
|
||||
self.use_zero_conv = False
|
||||
for i in range(len(self.channels)):
|
||||
if self.fusion_type == 'SPADE':
|
||||
# Corresponding to SPADE <Semantic Image Synthesis with Spatially-Adaptive Normalization>
|
||||
self.gamma.append(nn.Conv2d(out_channels[i], out_channels[i], 1, padding=0))
|
||||
self.beta.append(nn.Conv2d(out_channels[i], out_channels[i], 1, padding=0))
|
||||
self.norm.append(nn.BatchNorm2d(out_channels[i]))
|
||||
elif use_zero_conv:
|
||||
self.out.append(self.make_zero_conv(out_channels[i]))
|
||||
else:
|
||||
self.out.append(nn.Conv2d(out_channels[i], out_channels[i], 1, padding=0))
|
||||
for j in range(nums_rb):
|
||||
if i==0:
|
||||
# 1280, 32, 32 -> 1280, 32, 32
|
||||
self.body.append(
|
||||
ResnetBlock(in_channels[i], out_channels[i], down=False, up=False, ksize=ksize, sk=sk, use_conv=use_conv,
|
||||
enable_timestep=enable_timestep, temb_channels=temb_channels, use_norm=use_norm))
|
||||
# 1280, 32, 32 -> 1280, 32, 32
|
||||
elif i==1:
|
||||
# 640, 64, 64 -> 1280, 64, 64
|
||||
if j==0:
|
||||
self.body.append(
|
||||
ResnetBlock(in_channels[i], out_channels[i], down=False, up=False, ksize=ksize, sk=sk,
|
||||
use_conv=use_conv, enable_timestep=enable_timestep, temb_channels=temb_channels, use_norm=use_norm))
|
||||
else:
|
||||
self.body.append(
|
||||
ResnetBlock(out_channels[i], out_channels[i], down=False, up=False, ksize=ksize,sk=sk,
|
||||
use_conv=use_conv, enable_timestep=enable_timestep, temb_channels=temb_channels, use_norm=use_norm))
|
||||
else:
|
||||
# 320, 64, 64 -> 640, 128, 128
|
||||
if j==0:
|
||||
self.body.append(
|
||||
ResnetBlock(in_channels[i], out_channels[i], down=False, up=True, ksize=ksize, sk=sk,
|
||||
use_conv=True, enable_timestep=enable_timestep, temb_channels=temb_channels, use_norm=use_norm))
|
||||
# use convtranspose2d
|
||||
else:
|
||||
self.body.append(
|
||||
ResnetBlock(out_channels[i], out_channels[i], down=False, up=False, ksize=ksize, sk=sk,
|
||||
use_conv=use_conv, enable_timestep=enable_timestep, temb_channels=temb_channels, use_norm=use_norm))
|
||||
|
||||
|
||||
self.body = nn.ModuleList(self.body)
|
||||
if self.use_zero_conv:
|
||||
self.zero_out = nn.ModuleList(self.out)
|
||||
|
||||
# if self.fusion_type == 'SPADE':
|
||||
# self.norm = nn.ModuleList(self.norm)
|
||||
# self.gamma = nn.ModuleList(self.gamma)
|
||||
# self.beta = nn.ModuleList(self.beta)
|
||||
# else:
|
||||
# self.zero_out = nn.ModuleList(self.out)
|
||||
|
||||
|
||||
# if enable_timestep:
|
||||
# a = 320
|
||||
#
|
||||
# time_embed_dim = a * 4
|
||||
# self.time_proj = Timesteps(a, True, 0)
|
||||
# timestep_input_dim = a
|
||||
#
|
||||
# self.time_embedding = TimestepEmbedding(
|
||||
# timestep_input_dim,
|
||||
# time_embed_dim,
|
||||
# act_fn='silu',
|
||||
# post_act_fn=None,
|
||||
# cond_proj_dim=None,
|
||||
# )
|
||||
|
||||
|
||||
def make_zero_conv(self, channels):
|
||||
|
||||
return zero_module(nn.Conv2d(channels, channels, 1, padding=0))
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
"""
|
||||
`torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
|
||||
"""
|
||||
return get_parameter_dtype(self)
|
||||
|
||||
def forward(self, x, t=None):
|
||||
# extract features
|
||||
features = []
|
||||
b, c, _, _ = x[-1].shape
|
||||
if t is not None:
|
||||
if not torch.is_tensor(t):
|
||||
is_mps = x[0].device.type == "mps"
|
||||
if isinstance(timestep, float):
|
||||
dtype = torch.float32 if is_mps else torch.float64
|
||||
else:
|
||||
dtype = torch.int32 if is_mps else torch.int64
|
||||
t = torch.tensor([t], dtype=dtype, device=x[0].device)
|
||||
elif len(t.shape) == 0:
|
||||
t = t[None].to(x[0].device)
|
||||
|
||||
t = t.expand(b)
|
||||
t = self.time_proj(t) # b, 320
|
||||
t = t.to(dtype=x[0].dtype)
|
||||
t = self.time_embedding(t) # b, 1280
|
||||
output_size = (b, 640, 128, 128) # last CA layer output
|
||||
for i in range(len(self.channels)):
|
||||
for j in range(self.nums_rb):
|
||||
idx = i * self.nums_rb + j
|
||||
if j == 0:
|
||||
if i < 2:
|
||||
out = self.body[idx](x[i], temb=t)
|
||||
else:
|
||||
out = self.body[idx](x[i], output_size=output_size, temb=t)
|
||||
else:
|
||||
out = self.body[idx](out, temb=t)
|
||||
if self.fusion_type == 'SPADE':
|
||||
out_gamma = self.gamma[i](out)
|
||||
out_beta = self.beta[i](out)
|
||||
out = [out_gamma, out_beta]
|
||||
else:
|
||||
out = self.zero_out[i](out)
|
||||
features.append(out)
|
||||
|
||||
return features
|
||||
|
||||
|
||||
def zero_module(module):
|
||||
"""
|
||||
Zero out the parameters of a module and return it.
|
||||
"""
|
||||
for p in module.parameters():
|
||||
p.detach().zero_()
|
||||
return module
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,14 @@
|
||||
import os
|
||||
import imageio
|
||||
import numpy as np
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
import torchvision
|
||||
import torch.distributed as dist
|
||||
|
||||
from safetensors import safe_open
|
||||
from tqdm import tqdm
|
||||
from einops import rearrange
|
||||
from model.convert_from_ckpt import convert_ldm_unet_checkpoint, convert_ldm_clip_checkpoint, convert_ldm_vae_checkpoint
|
||||
# from animatediff.utils.convert_lora_safetensor_to_diffusers import convert_lora, convert_motion_lora_ckpt_to_diffusers
|
||||
@@ -0,0 +1,58 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class FourierEmbedder(nn.Module):
|
||||
def __init__(self, num_freqs=64, temperature=100):
|
||||
super().__init__()
|
||||
|
||||
self.num_freqs = num_freqs
|
||||
self.temperature = temperature
|
||||
|
||||
freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
|
||||
freq_bands = freq_bands[None, None, None]
|
||||
self.register_buffer("freq_bands", freq_bands, persistent=False)
|
||||
|
||||
def __call__(self, x):
|
||||
x = self.freq_bands * x.unsqueeze(-1)
|
||||
return torch.stack((x.sin(), x.cos()), dim=-1).permute(0, 1, 3, 4, 2).reshape(*x.shape[:2], -1)
|
||||
|
||||
|
||||
class PositionNet(nn.Module):
|
||||
def __init__(self, positive_len, out_dim, fourier_freqs=8):
|
||||
super().__init__()
|
||||
self.positive_len = positive_len
|
||||
self.out_dim = out_dim
|
||||
|
||||
self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
|
||||
self.position_dim = fourier_freqs * 2 * 4 # 2: sin/cos, 4: xyxy
|
||||
|
||||
if isinstance(out_dim, tuple):
|
||||
out_dim = out_dim[0]
|
||||
self.linears = nn.Sequential(
|
||||
nn.Linear(self.positive_len + self.position_dim, 512),
|
||||
nn.SiLU(),
|
||||
nn.Linear(512, 512),
|
||||
nn.SiLU(),
|
||||
nn.Linear(512, out_dim),
|
||||
)
|
||||
|
||||
self.null_positive_feature = torch.nn.Parameter(torch.zeros([self.positive_len]))
|
||||
self.null_position_feature = torch.nn.Parameter(torch.zeros([self.position_dim]))
|
||||
|
||||
def forward(self, boxes, masks, positive_embeddings):
|
||||
masks = masks.unsqueeze(-1)
|
||||
|
||||
# embedding position (it may includes padding as placeholder)
|
||||
xyxy_embedding = self.fourier_embedder(boxes) # B*N*4 -> B*N*C
|
||||
|
||||
# learnable null embedding
|
||||
positive_null = self.null_positive_feature.view(1, 1, -1)
|
||||
xyxy_null = self.null_position_feature.view(1, 1, -1)
|
||||
|
||||
# replace padding with learnable null embedding
|
||||
positive_embeddings = positive_embeddings * masks + (1 - masks) * positive_null
|
||||
xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null
|
||||
|
||||
objs = self.linears(torch.cat([positive_embeddings, xyxy_embedding], dim=-1))
|
||||
return objs
|
||||
Reference in New Issue
Block a user