Spaces:
Running on Zero
Running on Zero
File size: 10,917 Bytes
f0395ef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 | # Copyright (c) 2025 Hansheng Chen
import sys
import inspect
import torch
import torch.nn as nn
import mmcv
import diffusers
from copy import deepcopy
from mmcv.runner.fp16_utils import force_fp32
from mmgen.models.architectures.common import get_module_device
from mmgen.models.builder import MODULES, build_module
from . import schedulers
@torch.jit.script
def guidance_jit(pos_mean, neg_mean, guidance_scale: float, orthogonal: bool = False):
bias = (pos_mean - neg_mean) * (guidance_scale - 1)
if orthogonal:
dim = list(range(1, pos_mean.dim()))
bias = bias - (bias * pos_mean).mean(
dim=dim, keepdim=True
) / (pos_mean * pos_mean).mean(dim=dim, keepdim=True).clamp(min=1e-6) * pos_mean
return bias
@MODULES.register_module()
class GaussianFlow(nn.Module):
def __init__(self,
denoising=None,
flow_loss=None,
num_timesteps=1000,
timestep_sampler=dict(type='ContinuousTimeStepSampler', shift=1.0),
flip_model_timesteps=False,
denoising_mean_mode='U',
train_cfg=None,
test_cfg=None):
super().__init__()
# build denoising module in this function
self.num_timesteps = num_timesteps
self.denoising = build_module(denoising) if isinstance(denoising, dict) else denoising
self.denoising_mean_mode = denoising_mean_mode
self.flip_model_timesteps = flip_model_timesteps
self.train_cfg = deepcopy(train_cfg) if train_cfg is not None else dict()
self.test_cfg = deepcopy(test_cfg) if test_cfg is not None else dict()
# build sampler
self.timestep_sampler = build_module(
timestep_sampler,
default_args=dict(num_timesteps=num_timesteps))
self.flow_loss = build_module(flow_loss) if flow_loss is not None else None
def forward_transition(
self, x_t_src, t_src=None, t_tgt=None, sigma_src=None, sigma_tgt=None, eps=1e-6):
if sigma_src is None:
if not isinstance(t_src, torch.Tensor):
t_src = torch.tensor(t_src, device=x_t_src.device)
t_src = t_src.reshape(*t_src.size(), *((x_t_src.dim() - t_src.dim()) * [1]))
sigma_src = t_src / self.num_timesteps
if sigma_tgt is None:
if not isinstance(t_tgt, torch.Tensor):
t_tgt = torch.tensor(t_tgt, device=x_t_src.device)
t_tgt = t_tgt.reshape(*t_tgt.size(), *((x_t_src.dim() - t_tgt.dim()) * [1]))
sigma_tgt = t_tgt / self.num_timesteps
alpha_src = 1 - sigma_src
alpha_tgt = 1 - sigma_tgt
scale_trans = alpha_tgt / alpha_src.clamp(min=eps)
var_trans = sigma_tgt ** 2 - (scale_trans * sigma_src) ** 2
return dict(mean=x_t_src * scale_trans, var=var_trans), scale_trans
def sample_forward_transition(self, x_t_src, noise, t_src=None, t_tgt=None, sigma_src=None, sigma_tgt=None):
trans_g = self.forward_transition(
x_t_src, t_src=t_src, t_tgt=t_tgt, sigma_src=sigma_src, sigma_tgt=sigma_tgt)[0]
return trans_g['mean'] + noise * trans_g['var'].sqrt()
def sample_forward_diffusion(self, x_0, t, noise):
if t.dim() == 0:
t = t.expand(x_0.size(0))
std = t.reshape(*t.size(), *((x_0.dim() - t.dim()) * [1])) / self.num_timesteps
mean = 1 - std
return x_0 * mean + noise * std, mean, std
def pred(self, x_t=None, t=None, **kwargs):
ori_dtype = x_t.dtype
if hasattr(self.denoising, 'dtype'):
denoising_dtype = self.denoising.dtype
else:
denoising_dtype = next(self.denoising.parameters()).dtype
x_t = x_t.to(denoising_dtype)
num_batches = x_t.size(0)
if t.dim() == 0 or len(t) != num_batches:
t = t.expand(num_batches)
if self.flip_model_timesteps:
t = self.num_timesteps - t
output = self.denoising(x_t, t, **kwargs)
if isinstance(output, dict):
output = {k: v.to(ori_dtype) for k, v in output.items()}
else:
output = output.to(ori_dtype)
return output
@force_fp32()
def loss(self, denoising_output, x_0, noise, t, pred_mask=None):
if self.denoising_mean_mode.upper() == 'U':
if isinstance(denoising_output, dict):
loss_kwargs = denoising_output
elif isinstance(denoising_output, torch.Tensor):
loss_kwargs = dict(u_t_pred=denoising_output)
else:
raise AttributeError('Unknown denoising output type '
f'[{type(denoising_output)}].')
loss_kwargs.update(u_t=noise - x_0)
else:
raise AttributeError('Unknown denoising mean output type '
f'[{self.denoising_mean_mode}].')
loss_kwargs.update(
x_0=x_0,
noise=noise,
timesteps=t,
weight=pred_mask.float() if pred_mask is not None else None)
return self.flow_loss(loss_kwargs)
def forward_train(self, x_0, **kwargs):
device = get_module_device(self)
num_batches = x_0.size(0)
seq_len = x_0.shape[2:].numel() # h * w or t * h * w
t = self.timestep_sampler(num_batches, seq_len=seq_len, device=device)
noise = torch.randn_like(x_0)
x_t, _, _ = self.sample_forward_diffusion(x_0, t, noise)
denoising_output = self.pred(x_t, t, **kwargs)
loss = self.loss(denoising_output, x_0, noise, t)
log_vars = self.flow_loss.log_vars
log_vars.update(loss_diffusion=float(loss))
return loss, log_vars
def forward_test(
self, x_0=None, noise=None, guidance_scale=1.0,
test_cfg_override=dict(), show_pbar=False, **kwargs):
x_t = torch.randn_like(x_0) if noise is None else noise
num_batches = x_t.size(0)
ori_dtype = x_t.dtype
x_t = x_t.float()
cfg = deepcopy(self.test_cfg)
cfg.update(test_cfg_override)
sampler = cfg.get('sampler', 'FlowEulerODE')
sampler_class = getattr(diffusers.schedulers, sampler + 'Scheduler', None)
if sampler_class is None:
sampler_class = getattr(schedulers, sampler + 'Scheduler', None)
if sampler_class is None:
raise AttributeError(f'Cannot find sampler [{sampler}].')
sampler_kwargs = cfg.get('sampler_kwargs', {})
signatures = inspect.signature(sampler_class).parameters.keys()
for key in ['shift', 'use_dynamic_shifting', 'base_seq_len', 'max_seq_len', 'base_logshift', 'max_logshift']:
if key in signatures and key not in sampler_kwargs:
sampler_kwargs[key] = cfg.get(key, getattr(self.timestep_sampler, key))
if 'flow_shift' in signatures and 'use_flow_sigmas' in signatures:
sampler_kwargs['prediction_type'] = 'flow_prediction'
sampler_kwargs['use_flow_sigmas'] = True
if 'flow_shift' not in sampler_kwargs:
sampler_kwargs['flow_shift'] = cfg.get('shift', self.timestep_sampler.shift)
sampler = sampler_class(self.num_timesteps, **sampler_kwargs)
num_timesteps = cfg.get('num_timesteps', self.num_timesteps)
guidance_interval = cfg.get('guidance_interval', [0, self.num_timesteps])
orthogonal_guidance = cfg.get('orthogonal_guidance', False)
use_guidance = guidance_scale > 1.0
set_timesteps_signatures = inspect.signature(sampler.set_timesteps).parameters.keys()
if 'seq_len' in set_timesteps_signatures:
seq_len = x_t.shape[2:].numel() # h * w or t * h * w
sampler.set_timesteps(num_timesteps, seq_len=seq_len, device=x_t.device)
else:
sampler.set_timesteps(num_timesteps, device=x_t.device)
timesteps = sampler.timesteps
if show_pbar:
pbar = mmcv.ProgressBar(len(timesteps))
for t in timesteps:
x_t_input = x_t
_kwargs = kwargs
if use_guidance:
guidance_active = guidance_interval[0] <= t <= guidance_interval[1]
if guidance_active:
x_t_input = torch.cat([x_t_input, x_t_input], dim=0)
else:
_kwargs = {
k: v[num_batches:] if isinstance(v, torch.Tensor) and v.size(0) == 2 * num_batches else v
for k, v in kwargs.items()}
denoising_output = self.pred(x_t_input, t, **_kwargs)
if use_guidance and guidance_active:
mean_neg, mean_pos = denoising_output.chunk(2, dim=0)
bias = guidance_jit(mean_pos, mean_neg, guidance_scale, orthogonal_guidance)
denoising_output = mean_pos + bias
x_t = sampler.step(denoising_output, t, x_t, return_dict=False)[0]
if show_pbar:
pbar.update()
if show_pbar:
sys.stdout.write('\n')
return x_t.to(ori_dtype)
def forward_u(self, x_t=None, t=None, guidance_scale=1.0, test_cfg_override=dict(), **kwargs):
ori_dtype = x_t.dtype
x_t = x_t.float()
num_batches = x_t.size(0)
cfg = deepcopy(self.test_cfg)
cfg.update(test_cfg_override)
orthogonal_guidance = cfg.get('orthogonal_guidance', False)
guidance_interval = cfg.get('guidance_interval', [0, self.num_timesteps])
use_guidance = guidance_scale > 1.0
x_t_input = x_t
t_input = t
if use_guidance:
x_t_input = torch.cat([x_t_input, x_t_input], dim=0)
t_input = torch.cat([t_input, t_input], dim=0)
denoising_output = self.pred(x_t_input, t_input, **kwargs)
if use_guidance:
mean_neg, mean_pos = denoising_output.chunk(2, dim=0)
bias = guidance_jit(mean_pos, mean_neg, guidance_scale, orthogonal_guidance)
if guidance_interval[0] > 0 or guidance_interval[1] < self.num_timesteps:
guidance_active = ((t >= guidance_interval[0]) & (t <= guidance_interval[1])).reshape(
[num_batches] + [1] * (bias.dim() - 1))
bias = bias.masked_fill(~guidance_active, 0.0)
denoising_output = mean_pos + bias
return denoising_output.to(ori_dtype)
def forward(
self,
x_0=None,
return_loss=False,
return_u=False,
return_denoising_output=False,
**kwargs):
if return_loss:
return self.forward_train(x_0, **kwargs)
elif return_u:
return self.forward_u(**kwargs)
elif return_denoising_output:
return self.pred(**kwargs)
else:
return self.forward_test(x_0, **kwargs)
|