from typing import Any, List, Optional, Union, Dict import torch from torch.optim.lr_scheduler import LambdaLR from safetensors.torch import load_file as load_safetensors from ..sgm.modules.diffusionmodules.openaimodel import * from ..sgm.modules.video_attention import SpatialVideoTransformer from ..sgm.util import default, instantiate_from_config, get_obj_from_str from ..sgm.modules.diffusionmodules.util import zero_module from ..sgm.modules.diffusionmodules.video_model import VideoResBlock, VideoUNet from ..sgm.models.diffusion import DiffusionEngine from ..sgm.util import instantiate_from_config, get_obj_from_str import comfy.ops ops = comfy.ops.manual_cast class ControlledVideoUNet(VideoUNet): def __init__( self, in_channels: int, model_channels: int, out_channels: int, num_res_blocks: int, attention_resolutions: int, dropout: float = 0.0, channel_mult: List[int] = (1, 2, 4, 8), conv_resample: bool = True, dims: int = 2, num_classes: Optional[int] = None, use_checkpoint: bool = False, num_heads: int = -1, num_head_channels: int = -1, num_heads_upsample: int = -1, use_scale_shift_norm: bool = False, resblock_updown: bool = False, transformer_depth: Union[List[int], int] = 1, transformer_depth_middle: Optional[int] = None, context_dim: Optional[int] = None, time_downup: bool = False, time_context_dim: Optional[int] = None, extra_ff_mix_layer: bool = False, use_spatial_context: bool = False, merge_strategy: str = "fixed", merge_factor: float = 0.5, spatial_transformer_attn_type: str = "softmax", temporal_attn_type = None, spatial_self_attn_type: str = None, conv3d_type: str = None, video_kernel_size: Union[int, List[int]] = 3, use_linear_in_transformer: bool = False, adm_in_channels: Optional[int] = None, disable_temporal_crossattention: bool = False, max_ddpm_temb_period: int = 10000, trainable_layers = None, ): nn.Module.__init__(self) assert context_dim is not None if num_heads_upsample == -1: num_heads_upsample = num_heads if num_heads == -1: assert num_head_channels != -1 if num_head_channels == -1: assert num_heads != -1 self.trainable_layers = trainable_layers self.in_channels = in_channels self.model_channels = model_channels self.out_channels = out_channels if isinstance(transformer_depth, int): transformer_depth = len(channel_mult) * [transformer_depth] transformer_depth_middle = default( transformer_depth_middle, transformer_depth[-1] ) self.num_res_blocks = num_res_blocks self.attention_resolutions = attention_resolutions self.dropout = dropout self.channel_mult = channel_mult self.conv_resample = conv_resample self.num_classes = num_classes self.use_checkpoint = use_checkpoint self.num_heads = num_heads self.num_head_channels = num_head_channels self.num_heads_upsample = num_heads_upsample self.temporal_attn_type = temporal_attn_type self.spatial_self_attn_type = spatial_self_attn_type self.conv3d_type = conv3d_type time_embed_dim = model_channels * 4 self.time_embed = nn.Sequential( linear(model_channels, time_embed_dim), nn.SiLU(), linear(time_embed_dim, time_embed_dim), ) if self.num_classes is not None: if isinstance(self.num_classes, int): self.label_emb = nn.Embedding(num_classes, time_embed_dim) elif self.num_classes == "continuous": print("setting up linear c_adm embedding layer") self.label_emb = ops.Conv2d(1, time_embed_dim) elif self.num_classes == "timestep": self.label_emb = nn.Sequential( Timestep(model_channels), nn.Sequential( linear(model_channels, time_embed_dim), nn.SiLU(), linear(time_embed_dim, time_embed_dim), ), ) elif self.num_classes == "sequential": assert adm_in_channels is not None self.label_emb = nn.Sequential( nn.Sequential( linear(adm_in_channels, time_embed_dim), nn.SiLU(), linear(time_embed_dim, time_embed_dim), ) ) else: raise ValueError() self.input_blocks = nn.ModuleList( [ TimestepEmbedSequential( conv_nd(dims, in_channels, model_channels, 3, padding=1) ) ] ) self._feature_size = model_channels input_block_chans = [model_channels] ch = model_channels ds = 1 def get_attention_layer( ch, num_heads, dim_head, depth=1, context_dim=None, use_checkpoint=False, disabled_sa=False, ): return SpatialVideoTransformer( ch, num_heads, dim_head, depth=depth, context_dim=context_dim, time_context_dim=time_context_dim, dropout=dropout, ff_in=extra_ff_mix_layer, use_spatial_context=use_spatial_context, merge_strategy=merge_strategy, merge_factor=merge_factor, checkpoint=use_checkpoint, use_linear=use_linear_in_transformer, attn_mode=spatial_transformer_attn_type, temporal_attn_mode=get_obj_from_str(temporal_attn_type) \ if temporal_attn_type is not None else None, disable_self_attn=disabled_sa, disable_temporal_crossattention=disable_temporal_crossattention, max_time_embed_period=max_ddpm_temb_period, spatial_self_attn_type=get_obj_from_str(spatial_self_attn_type) \ if spatial_self_attn_type is not None else None, ) def get_resblock( merge_factor, merge_strategy, video_kernel_size, ch, time_embed_dim, dropout, out_ch, dims, use_checkpoint, use_scale_shift_norm, down=False, up=False, ): return VideoResBlock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, channels=ch, emb_channels=time_embed_dim, dropout=dropout, out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, down=down, up=up, temporal_conv=get_obj_from_str(conv3d_type) \ if conv3d_type is not None else None, ) for level, mult in enumerate(channel_mult): for _ in range(num_res_blocks): layers = [ get_resblock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, ch=ch, time_embed_dim=time_embed_dim, dropout=dropout, out_ch=mult * model_channels, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ) ] ch = mult * model_channels if ds in attention_resolutions: if num_head_channels == -1: dim_head = ch // num_heads else: num_heads = ch // num_head_channels dim_head = num_head_channels layers.append( get_attention_layer( ch, num_heads, dim_head, depth=transformer_depth[level], context_dim=context_dim, use_checkpoint=use_checkpoint, disabled_sa=False, ) ) self.input_blocks.append(TimestepEmbedSequential(*layers)) self._feature_size += ch input_block_chans.append(ch) if level != len(channel_mult) - 1: ds *= 2 out_ch = ch self.input_blocks.append( TimestepEmbedSequential( get_resblock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, ch=ch, time_embed_dim=time_embed_dim, dropout=dropout, out_ch=out_ch, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, down=True, ) if resblock_updown else Downsample( ch, conv_resample, dims=dims, out_channels=out_ch, third_down=time_downup, ) ) ) ch = out_ch input_block_chans.append(ch) self._feature_size += ch if num_head_channels == -1: dim_head = ch // num_heads else: num_heads = ch // num_head_channels dim_head = num_head_channels self.middle_block = TimestepEmbedSequential( get_resblock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, ch=ch, time_embed_dim=time_embed_dim, out_ch=None, dropout=dropout, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ), get_attention_layer( ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim, use_checkpoint=use_checkpoint, ), get_resblock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, ch=ch, out_ch=None, time_embed_dim=time_embed_dim, dropout=dropout, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ), ) self._feature_size += ch self.output_blocks = nn.ModuleList([]) for level, mult in list(enumerate(channel_mult))[::-1]: for i in range(num_res_blocks + 1): ich = input_block_chans.pop() layers = [ get_resblock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, ch=ch + ich, time_embed_dim=time_embed_dim, dropout=dropout, out_ch=model_channels * mult, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ) ] ch = model_channels * mult if ds in attention_resolutions: if num_head_channels == -1: dim_head = ch // num_heads else: num_heads = ch // num_head_channels dim_head = num_head_channels layers.append( get_attention_layer( ch, num_heads, dim_head, depth=transformer_depth[level], context_dim=context_dim, use_checkpoint=use_checkpoint, disabled_sa=False, ) ) if level and i == num_res_blocks: out_ch = ch ds //= 2 layers.append( get_resblock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, ch=ch, time_embed_dim=time_embed_dim, dropout=dropout, out_ch=out_ch, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, up=True, ) if resblock_updown else Upsample( ch, conv_resample, dims=dims, out_channels=out_ch, third_up=time_downup, ) ) self.output_blocks.append(TimestepEmbedSequential(*layers)) self._feature_size += ch self.out = nn.Sequential( normalization(ch), nn.SiLU(), zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)), ) def forward( self, x: th.Tensor, timesteps: th.Tensor, context: Optional[th.Tensor] = None, y: Optional[th.Tensor] = None, time_context: Optional[th.Tensor] = None, control: Optional[th.Tensor] = None, num_video_frames: Optional[int] = None, image_only_indicator: Optional[th.Tensor] = None, ): assert (y is not None) == ( self.num_classes is not None ), "must specify y if and only if the model is class-conditional -> no, relax this TODO" hs = [] t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False) emb = self.time_embed(t_emb) if self.num_classes is not None: assert y.shape[0] == x.shape[0] emb = emb + self.label_emb(y) h = x for module in self.input_blocks: h = module( h, emb, context=context, image_only_indicator=image_only_indicator, time_context=time_context, num_video_frames=num_video_frames, ) hs.append(h) h = self.middle_block( h, emb, context=context, image_only_indicator=image_only_indicator, time_context=time_context, num_video_frames=num_video_frames, ) if control is not None: h += control.pop() for module in self.output_blocks: if control is None: h = th.cat([h, hs.pop()], dim=1) else: h = th.cat([h, hs.pop() + control.pop()], dim=1) h = module( h, emb, context=context, image_only_indicator=image_only_indicator, time_context=time_context, num_video_frames=num_video_frames, ) h = h.type(x.dtype) return self.out(h) def set_parameters_requires_grad(self): self.requires_grad_(False) if self.trainable_layers is not None: for name, module in self.named_modules(): if module.__class__.__name__ in self.trainable_layers: module.requires_grad_(True) a = 0 def get_trainable_parameters(self): params = [] for p in self.parameters(): if p.requires_grad: params += [p] return params def get_state_dict(self): dicts = [] for name, param in self.named_parameters(): if param.requires_grad: dicts.append( f'model.diffusion_model.{name}' ) return dicts class ControlNet(nn.Module): def __init__( self, in_channels: int, model_channels: int, hint_channels: int, num_res_blocks: int, attention_resolutions: int, dropout: float = 0.0, channel_mult: List[int] = (1, 2, 4, 8), conv_resample: bool = True, dims: int = 2, num_classes: Optional[int] = None, use_checkpoint: bool = False, num_heads: int = -1, num_head_channels: int = -1, num_heads_upsample: int = -1, use_scale_shift_norm: bool = False, resblock_updown: bool = False, transformer_depth: Union[List[int], int] = 1, transformer_depth_middle: Optional[int] = None, context_dim: Optional[int] = None, time_downup: bool = False, time_context_dim: Optional[int] = None, extra_ff_mix_layer: bool = False, use_spatial_context: bool = False, merge_strategy: str = "fixed", merge_factor: float = 0.5, spatial_transformer_attn_type: str = "softmax", temporal_attn_type: str = None, spatial_self_attn_type: str = None, conv3d_type: str = None, video_kernel_size: Union[int, List[int]] = 3, use_linear_in_transformer: bool = False, adm_in_channels: Optional[int] = None, disable_temporal_crossattention: bool = False, max_ddpm_temb_period: int = 10000, ): super().__init__() assert context_dim is not None if num_heads_upsample == -1: num_heads_upsample = num_heads if num_heads == -1: assert num_head_channels != -1 if num_head_channels == -1: assert num_heads != -1 self.dims = dims self.in_channels = in_channels self.model_channels = model_channels self.hint_channels = hint_channels if isinstance(transformer_depth, int): transformer_depth = len(channel_mult) * [transformer_depth] transformer_depth_middle = default( transformer_depth_middle, transformer_depth[-1] ) self.num_res_blocks = num_res_blocks self.attention_resolutions = attention_resolutions self.dropout = dropout self.channel_mult = channel_mult self.conv_resample = conv_resample self.num_classes = num_classes self.use_checkpoint = use_checkpoint self.num_heads = num_heads self.num_head_channels = num_head_channels self.num_heads_upsample = num_heads_upsample self.context_dim = context_dim self.adm_in_channels = adm_in_channels self.temporal_attn_type = temporal_attn_type self.spatial_self_attn_type = spatial_self_attn_type self.conv3d_type = conv3d_type time_embed_dim = model_channels * 4 self.time_embed = nn.Sequential( linear(model_channels, time_embed_dim), nn.SiLU(), linear(time_embed_dim, time_embed_dim), ) if self.num_classes is not None: if isinstance(self.num_classes, int): self.label_emb = nn.Embedding(num_classes, time_embed_dim) elif self.num_classes == "continuous": print("setting up linear c_adm embedding layer") self.label_emb = ops.Conv2d(1, time_embed_dim) elif self.num_classes == "timestep": self.label_emb = nn.Sequential( Timestep(model_channels), nn.Sequential( linear(model_channels, time_embed_dim), nn.SiLU(), linear(time_embed_dim, time_embed_dim), ), ) elif self.num_classes == "sequential": assert adm_in_channels is not None self.label_emb = nn.Sequential( nn.Sequential( linear(adm_in_channels, time_embed_dim), nn.SiLU(), linear(time_embed_dim, time_embed_dim), ) ) else: raise ValueError() self.input_blocks = nn.ModuleList( [ TimestepEmbedSequential( conv_nd(dims, in_channels, model_channels, 3, padding=1) ) ] ) self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels)]) self.input_hint_block = TimestepEmbedSequential( conv_nd(dims, hint_channels, 16, 3, padding=1), nn.SiLU(), conv_nd(dims, 16, 16, 3, padding=1), nn.SiLU(), conv_nd(dims, 16, 32, 3, padding=1, stride=2), nn.SiLU(), conv_nd(dims, 32, 32, 3, padding=1), nn.SiLU(), conv_nd(dims, 32, 96, 3, padding=1, stride=2), nn.SiLU(), conv_nd(dims, 96, 96, 3, padding=1), nn.SiLU(), conv_nd(dims, 96, 256, 3, padding=1, stride=2), nn.SiLU(), zero_module(conv_nd(dims, 256, model_channels, 3, padding=1)) ) self._feature_size = model_channels input_block_chans = [model_channels] ch = model_channels ds = 1 def get_attention_layer( ch, num_heads, dim_head, depth=1, context_dim=None, use_checkpoint=False, disabled_sa=False, ): return SpatialVideoTransformer( ch, num_heads, dim_head, depth=depth, context_dim=context_dim, time_context_dim=time_context_dim, dropout=dropout, ff_in=extra_ff_mix_layer, use_spatial_context=use_spatial_context, merge_strategy=merge_strategy, merge_factor=merge_factor, checkpoint=use_checkpoint, use_linear=use_linear_in_transformer, attn_mode=spatial_transformer_attn_type, temporal_attn_mode=get_obj_from_str(temporal_attn_type) \ if temporal_attn_type is not None else None, disable_self_attn=disabled_sa, disable_temporal_crossattention=disable_temporal_crossattention, max_time_embed_period=max_ddpm_temb_period, spatial_self_attn_type=get_obj_from_str(spatial_self_attn_type) \ if spatial_self_attn_type is not None else None, ) def get_resblock( merge_factor, merge_strategy, video_kernel_size, ch, time_embed_dim, dropout, out_ch, dims, use_checkpoint, use_scale_shift_norm, down=False, up=False, ): return VideoResBlock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, channels=ch, emb_channels=time_embed_dim, dropout=dropout, out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, down=down, up=up, temporal_conv=get_obj_from_str(conv3d_type) \ if conv3d_type is not None else None, ) for level, mult in enumerate(channel_mult): for _ in range(num_res_blocks): layers = [ get_resblock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, ch=ch, time_embed_dim=time_embed_dim, dropout=dropout, out_ch=mult * model_channels, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ) ] ch = mult * model_channels if ds in attention_resolutions: if num_head_channels == -1: dim_head = ch // num_heads else: num_heads = ch // num_head_channels dim_head = num_head_channels layers.append( get_attention_layer( ch, num_heads, dim_head, depth=transformer_depth[level], context_dim=context_dim, use_checkpoint=use_checkpoint, disabled_sa=False, ) ) self.input_blocks.append(TimestepEmbedSequential(*layers)) self.zero_convs.append(self.make_zero_conv(ch)) self._feature_size += ch input_block_chans.append(ch) if level != len(channel_mult) - 1: ds *= 2 out_ch = ch self.input_blocks.append( TimestepEmbedSequential( get_resblock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, ch=ch, time_embed_dim=time_embed_dim, dropout=dropout, out_ch=out_ch, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, down=True, ) if resblock_updown else Downsample( ch, conv_resample, dims=dims, out_channels=out_ch, third_down=time_downup, ) ) ) ch = out_ch input_block_chans.append(ch) self.zero_convs.append(self.make_zero_conv(ch)) self._feature_size += ch if num_head_channels == -1: dim_head = ch // num_heads else: num_heads = ch // num_head_channels dim_head = num_head_channels self.middle_block = TimestepEmbedSequential( get_resblock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, ch=ch, time_embed_dim=time_embed_dim, out_ch=None, dropout=dropout, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ), get_attention_layer( ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim, use_checkpoint=use_checkpoint, ), get_resblock( merge_factor=merge_factor, merge_strategy=merge_strategy, video_kernel_size=video_kernel_size, ch=ch, out_ch=None, time_embed_dim=time_embed_dim, dropout=dropout, dims=dims, use_checkpoint=use_checkpoint, use_scale_shift_norm=use_scale_shift_norm, ), ) self.middle_block_out = self.make_zero_conv(ch) self._feature_size += ch def make_zero_conv(self, channels): return TimestepEmbedSequential(zero_module(conv_nd(self.dims, channels, channels, 1, padding=0))) def forward( self, x: th.Tensor, hint: th.Tensor, timesteps: th.Tensor, context: Optional[th.Tensor] = None, y: Optional[th.Tensor] = None, time_context: Optional[th.Tensor] = None, num_video_frames: Optional[int] = None, image_only_indicator: Optional[th.Tensor] = None, ): assert (y is not None) == ( self.num_classes is not None ), "must specify y if and only if the model is class-conditional -> no, relax this TODO" t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False) emb = self.time_embed(t_emb) if self.num_classes is not None: assert y.shape[0] == x.shape[0] emb = emb + self.label_emb(y) guided_hint = self.input_hint_block(hint, emb, context) outs = [] h = x for module, zero_conv in zip(self.input_blocks, self.zero_convs): h = module( h, emb, context=context, image_only_indicator=image_only_indicator, time_context=time_context, num_video_frames=num_video_frames, ) if guided_hint is not None: h += guided_hint guided_hint = None outs.append(zero_conv(h, emb, context)) h = self.middle_block( h, emb, context=context, image_only_indicator=image_only_indicator, time_context=time_context, num_video_frames=num_video_frames, ) outs.append(self.middle_block_out(h, emb, context)) return outs def init_from_ckpt( self, path: str, ) -> None: if path.endswith("ckpt"): sd = torch.load(path, map_location="cpu")["state_dict"] elif path.endswith("safetensors"): sd = load_safetensors(path) else: raise NotImplementedError missing, unexpected = self.load_state_dict(sd, strict=False) #print( # f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys" #) # if len(missing) > 0: # print(f"Missing Keys: {missing}") # if len(unexpected) > 0: # print(f"Unexpected Keys: {unexpected}") def set_parameters_requires_grad(self): self.requires_grad_(True) def get_trainable_parameters(self): params = [] for p in self.parameters(): if p.requires_grad: params += [p] return params def get_state_dict(self): dicts = [] for name, param in self.named_parameters(): if param.requires_grad: dicts.append( f'control_model.{name}' ) return dicts class VideoDiffusionEngine(DiffusionEngine): def __init__( self, controlnet_config, control_model_path=None, init_from_unet=False, global_average_pooling=False, sd_locked=True, drop_first_stage_model=False, *args, **kwargs ): super().__init__(*args, **kwargs) self.sd_locked = sd_locked self.control_model = instantiate_from_config(controlnet_config) self.control_scales = [1.0] * 13 self.global_average_pooling = global_average_pooling self.conditioner.embedders[0].eval() self.conditioner.embedders[0].requires_grad_(False) self.conditioner.embedders[3].eval() self.conditioner.embedders[3].requires_grad_(False) self.first_stage_model.eval() self.first_stage_model.requires_grad_(False) if self.sd_locked: self.model.diffusion_model.eval() self.model.diffusion_model.requires_grad_(False) else: self.model.diffusion_model.train() self.model.diffusion_model.set_parameters_requires_grad() self.control_model.train() self.control_model.set_parameters_requires_grad() if init_from_unet: missing, unexpected = self.control_model.load_state_dict( self.model.diffusion_model.state_dict(), strict=False ) #print(f"Restored from UNet {len(missing)} missing and {len(unexpected)} unexpected keys") # Expand K, V layers for TripleAttention for name, module in self.named_modules(): if '.time_stack' not in name and name[-5:] == 'attn1' and hasattr(module, 'expand'): module.expand() if control_model_path is not None: self.init_from_ckpt(control_model_path) if drop_first_stage_model: del self.first_stage_model def forward(self, x, batch): loss = self.loss_fn(self.apply_model, self.denoiser, self.conditioner, x, batch) loss_mean = loss.mean() loss_dict = {"loss": loss_mean} return loss_mean, loss_dict def shared_step(self, batch: Dict) -> Any: x = self.get_input(batch) x.requires_grad_(True) if x.shape[1] == 3: x = self.encode_first_stage(x) batch["global_step"] = self.global_step loss, loss_dict = self(x, batch) return loss, loss_dict def apply_model( self, x: th.Tensor, timesteps: th.Tensor, cond: Dict, time_context: Optional[th.Tensor] = None, num_video_frames: Optional[int] = None, image_only_indicator: Optional[th.Tensor] = None, ): B, _, _, _ = x.shape context = cond.get('crossattn', None) if 'crossattn_scale' in cond.keys(): context = context * cond['crossattn_scale'] y = cond.get('vector', None) control_hint = cond.get('control_hint', None) cond_concat = cond.get("concat", torch.Tensor([])) cond_concat = cond_concat.type_as(x) if 'concat_scale' in cond.keys(): cond_concat = cond_concat * cond['concat_scale'] input_x = torch.cat([x, cond_concat], dim=1) if control_hint is not None: controls = self.control_model( x=input_x, hint=control_hint, timesteps=timesteps, context=context, y=y, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator, ) controls = [c * scale for c, scale in zip(controls, self.control_scales)] if self.global_average_pooling: controls = [torch.mean(c, dim=(2, 3), keepdim=True) for c in controls] else: controls = None out = self.model.diffusion_model( x=input_x, timesteps=timesteps, context=context, y=y, time_context=time_context, control=controls, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator, ) return out def configure_optimizers(self): lr = self.learning_rate params = self.control_model.get_trainable_parameters() if not self.sd_locked: params += self.model.diffusion_model.get_trainable_parameters() for embedder in self.conditioner.embedders: if embedder.is_trainable: params = params + list(embedder.parameters()) opt = self.instantiate_optimizer_from_config(params, lr, self.optimizer_config) if self.scheduler_config is not None: scheduler = instantiate_from_config(self.scheduler_config) print("Setting up LambdaLR scheduler...") scheduler = [ { "scheduler": LambdaLR(opt, lr_lambda=scheduler.schedule), "interval": "step", "frequency": 1, } ] return [opt], scheduler return opt def on_save_checkpoint(self, checkpoint): save_dict = self.control_model.get_state_dict() + self.model.diffusion_model.get_state_dict() keys = list( checkpoint['state_dict'].keys() ) for k in keys: if k not in save_dict: del checkpoint['state_dict'][k] def on_load_checkpoint(self, checkpoint): load_keys = list( checkpoint['state_dict'].keys() ) # Change learning rate when resume #checkpoint['optimizer_states'][0]['param_groups'][0]['lr'] = 1.0e-5 self_state_dict = self.state_dict() keys = list( self_state_dict.keys() ) for k in keys: if k not in load_keys: checkpoint['state_dict'][k] = self_state_dict[k]