From 92aae58f6e55a0b72c89a700a7f8f472421fc3d3 Mon Sep 17 00:00:00 2001 From: Kijai <40791699+kijai@users.noreply.github.com> Date: Wed, 14 Feb 2024 15:46:31 +0200 Subject: [PATCH] resent --- nodes.py | 26 +++++-- temp.py | 216 +------------------------------------------------------ 2 files changed, 22 insertions(+), 220 deletions(-) diff --git a/nodes.py b/nodes.py index 7d96d72..23e5b17 100644 --- a/nodes.py +++ b/nodes.py @@ -199,7 +199,7 @@ class ADMD_InitializeTraining: ], { "default": 'Lion' }), - + "include_resnet": ("BOOLEAN", {"default": False}), }, } @@ -212,7 +212,7 @@ class ADMD_InitializeTraining: def process(self, pipeline, images, prompt, lora_name, learning_rate, learning_rate_spatial, - lora_rank, seed, optimization_method, max_train_steps): + lora_rank, seed, optimization_method, max_train_steps, include_resnet): with torch.inference_mode(False): validation_pipeline = pipeline["validation_pipeline"] @@ -246,7 +246,11 @@ class ADMD_InitializeTraining: is_debug = False lora_unet_dropout = 0.1 - target_spatial_modules = ["Transformer3DModel"] + if include_resnet: + target_spatial_modules = ["Transformer3DModel", "ResnetBlock2D"] + else: + target_spatial_modules = ["Transformer3DModel"] + target_temporal_modules = ["TemporalTransformerBlock"] name = lora_name @@ -372,7 +376,8 @@ class ADMD_InitializeTraining: "train_noise_scheduler_spatial": train_noise_scheduler_spatial, "validation_pipeline": validation_pipeline, "global_step": 0, - "scaler": scaler + "scaler": scaler, + "include_resnet": include_resnet } #Data batch sanity check @@ -805,11 +810,13 @@ class ADMD_TrainLora: @classmethod def INPUT_TYPES(s): return { - "required": { - "trigger_input": ("VHS_FILENAMES", ), + "required": { "admd_pipeline": ("ADMDPIPELINE", ), "steps": ("INT", {"default": 100, "min": 0, "max": 10000, "step": 1}), }, + "optional": { + "trigger_input": ("VHS_FILENAMES", ), #attempt to force comfy execution order + } } RETURN_TYPES = ("ADMDPIPELINE",) @@ -833,6 +840,7 @@ class ADMD_TrainLora: pixel_values = admd_pipeline["pixel_values"] scaler = admd_pipeline["scaler"] + include_resnet = admd_pipeline["include_resnet"] use_offset_noise = False device = comfy.model_management.get_torch_device() @@ -847,7 +855,11 @@ class ADMD_TrainLora: unet.enable_gradient_checkpointing() unet.train() - target_spatial_modules = ["Transformer3DModel"] + if include_resnet: + target_spatial_modules = ["Transformer3DModel", "ResnetBlock2D"] + else: + target_spatial_modules = ["Transformer3DModel"] + target_temporal_modules = ["TemporalTransformerBlock"] batch_size = 1 diff --git a/temp.py b/temp.py index e866100..653a6fa 100644 --- a/temp.py +++ b/temp.py @@ -1,213 +1,3 @@ -AutoencoderKL( - (encoder): Encoder( - (conv_in): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (down_blocks): ModuleList( - (0): DownEncoderBlock2D( - (resnets): ModuleList( - (0-1): 2 x ResnetBlock2D( - (norm1): GroupNorm(32, 128, eps=1e-06, affine=True) - (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 128, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - ) - ) - (downsamplers): ModuleList( - (0): Downsample2D( - (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2)) - ) - ) - ) - (1): DownEncoderBlock2D( - (resnets): ModuleList( - (0): ResnetBlock2D( - (norm1): GroupNorm(32, 128, eps=1e-06, affine=True) - (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 256, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - (conv_shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1)) - ) - (1): ResnetBlock2D( - (norm1): GroupNorm(32, 256, eps=1e-06, affine=True) - (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 256, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - ) - ) - (downsamplers): ModuleList( - (0): Downsample2D( - (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2)) - ) - ) - ) - (2): DownEncoderBlock2D( - (resnets): ModuleList( - (0): ResnetBlock2D( - (norm1): GroupNorm(32, 256, eps=1e-06, affine=True) - (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 512, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - (conv_shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1)) - ) - (1): ResnetBlock2D( - (norm1): GroupNorm(32, 512, eps=1e-06, affine=True) - (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 512, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - ) - ) - (downsamplers): ModuleList( - (0): Downsample2D( - (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2)) - ) - ) - ) - (3): DownEncoderBlock2D( - (resnets): ModuleList( - (0-1): 2 x ResnetBlock2D( - (norm1): GroupNorm(32, 512, eps=1e-06, affine=True) - (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 512, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - ) - ) - ) - ) - (mid_block): UNetMidBlock2D( - (attentions): ModuleList( - (0): Attention( - (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True) - (to_q): Linear(in_features=512, out_features=512, bias=True) - (to_k): Linear(in_features=512, out_features=512, bias=True) - (to_v): Linear(in_features=512, out_features=512, bias=True) - (to_out): ModuleList( - (0): Linear(in_features=512, out_features=512, bias=True) - (1): Dropout(p=0.0, inplace=False) - ) - ) - ) - (resnets): ModuleList( - (0-1): 2 x ResnetBlock2D( - (norm1): GroupNorm(32, 512, eps=1e-06, affine=True) - (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 512, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - ) - ) - ) - (conv_norm_out): GroupNorm(32, 512, eps=1e-06, affine=True) - (conv_act): SiLU() - (conv_out): Conv2d(512, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - ) - (decoder): Decoder( - (conv_in): Conv2d(4, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (up_blocks): ModuleList( - (0-1): 2 x UpDecoderBlock2D( - (resnets): ModuleList( - (0-2): 3 x ResnetBlock2D( - (norm1): GroupNorm(32, 512, eps=1e-06, affine=True) - (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 512, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - ) - ) - (upsamplers): ModuleList( - (0): Upsample2D( - (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - ) - ) - ) - (2): UpDecoderBlock2D( - (resnets): ModuleList( - (0): ResnetBlock2D( - (norm1): GroupNorm(32, 512, eps=1e-06, affine=True) - (conv1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 256, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - (conv_shortcut): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1)) - ) - (1-2): 2 x ResnetBlock2D( - (norm1): GroupNorm(32, 256, eps=1e-06, affine=True) - (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 256, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - ) - ) - (upsamplers): ModuleList( - (0): Upsample2D( - (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - ) - ) - ) - (3): UpDecoderBlock2D( - (resnets): ModuleList( - (0): ResnetBlock2D( - (norm1): GroupNorm(32, 256, eps=1e-06, affine=True) - (conv1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 128, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - (conv_shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1)) - ) - (1-2): 2 x ResnetBlock2D( - (norm1): GroupNorm(32, 128, eps=1e-06, affine=True) - (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 128, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - ) - ) - ) - ) - (mid_block): UNetMidBlock2D( - (attentions): ModuleList( - (0): Attention( - (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True) - (to_q): Linear(in_features=512, out_features=512, bias=True) - (to_k): Linear(in_features=512, out_features=512, bias=True) - (to_v): Linear(in_features=512, out_features=512, bias=True) - (to_out): ModuleList( - (0): Linear(in_features=512, out_features=512, bias=True) - (1): Dropout(p=0.0, inplace=False) - ) - ) - ) - (resnets): ModuleList( - (0-1): 2 x ResnetBlock2D( - (norm1): GroupNorm(32, 512, eps=1e-06, affine=True) - (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (norm2): GroupNorm(32, 512, eps=1e-06, affine=True) - (dropout): Dropout(p=0.0, inplace=False) - (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - (nonlinearity): SiLU() - ) - ) - ) - (conv_norm_out): GroupNorm(32, 128, eps=1e-06, affine=True) - (conv_act): SiLU() - (conv_out): Conv2d(128, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - ) - (quant_conv): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1)) - (post_quant_conv): Conv2d(4, 4, kernel_size=(1, 1), stride=(1, 1)) -) +{'sample_size': 64, 'in_channels': 4, 'down_block_types': ('CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'DownBlock2D'), 'block_out_channels': (320, 640, 1280, 1280), 'layers_per_block': 2, 'cross_attention_dim': 768, 'attention_head_dim': 8, 'use_linear_projection': False, 'class_embed_type': None, 'addition_embed_type': None, 'addition_time_embed_dim': None, 'projection_class_embeddings_input_dim': None, 'transformer_layers_per_block': 1, 'out_channels': 4, 'up_block_types': ('UpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D')} +'sample_size', 64), ('in_channels', 4), ('out_channels', 4), ('center_input_sample', False), ('flip_sin_to_cos', True), ('freq_shift', 0), ('down_block_types', ['CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'DownBlock3D']), ('mid_block_type', 'UNetMidBlock3DCrossAttn'), ('up_block_types', ['UpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D']), ('only_cross_attention', False), ('block_out_channels', [320, 640, 1280, 1280]), ('layers_per_block', 2), ('downsample_padding', 1), ('mid_block_scale_factor', 1), ('act_fn', 'silu'), ('norm_num_groups', 32), ('norm_eps', 1e-05), ('cross_attention_dim', 768), ('attention_head_dim', 8), ('dual_cross_attention', False), ('use_linear_projection', False), ('class_embed_type', None), ('num_class_embeds', None), ('upcast_attention', False), ('resnet_time_scale_shift', 'default'), ('use_inflated_groupnorm', True), ('use_motion_module', True), ('motion_module_resolutions', [1, 2, 4, 8]), ('motion_module_mid_block', False), ('motion_module_decoder_only', False), ('motion_module_type', 'Vanilla'), ('motion_module_kwargs', {'num_attention_heads': 8, 'num_transformer_block': 1, 'attention_block_types': ['Temporal_Self', 'Temporal_Self'], 'temporal_position_encoding': True, 'temporal_position_encoding_max_len': 32, 'temporal_attention_dim_div': 1, 'zero_initialize': True}), ('unet_use_cross_frame_attention', False), ('unet_use_temporal_attention', False), ('_use_default_values', ['resnet_time_scale_shift', 'only_cross_attention', 'mid_block_type', 'unet_use_cross_frame_attention', 'class_embed_type', 'unet_use_temporal_attention', 'dual_cross_attention', 'num_class_embeds', 'upcast_attention', 'use_linear_projection', 'motion_module_decoder_only']), ('_class_name', 'UNet3DConditionModel'), ('_diffusers_version', '0.6.0')]) + \ No newline at end of file