Files: Epoch -1

config.py

@@ -0,0 +1,22 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class TrainingConfig:
+    image_size = 128  # the generated image resolution
+    train_batch_size = 4
+    eval_batch_size = 4  # how many images to sample during evaluation
+    num_epochs = 50
+    gradient_accumulation_steps = 1
+    learning_rate = 1e-4
+    lr_warmup_steps = 500
+    save_image_epochs = 1
+    save_model_epochs = 3
+    mixed_precision = 'fp16'  # `no` for float32, `fp16` for automatic mixed precision
+    output_dir = 'ddpm-paintings-128-finetuned-cifar10'  # the model name locally and on the HF Hub
+
+    push_to_hub = True  # whether to upload the saved model to the HF Hub
+    hub_model_id = 'jmemon/ddpm-paintings-128-finetuned-cifar10'  # the name of the repository to create on the HF Hub
+    hub_private_repo = False
+    overwrite_output_dir = True  # overwrite the old model when re-running the notebook
+    seed = 0

ddpm-paintings-128-finetuned-cifar10/logs/ddpm-paintings-128-finetuned-cifar10/events.out.tfevents.1701696166.coffee.14798.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2fbf486914eb9ed63fdbcf637c2874ca608a32f1ec948a4567e37a8e2e412f3
+size 427942

ddpm-paintings-128-finetuned-cifar10/logs/ddpm-paintings-128-finetuned-cifar10/events.out.tfevents.1701704512.coffee.17529.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c99374f9a97092f9da24ebc79289a21d0e48e40598c677c8206b1f453c2b050
+size 88

main.py

@@ -0,0 +1,200 @@
+from pathlib import Path
+import PIL
+from tqdm import tqdm
+
+from accelerate import Accelerator
+from datasets import load_dataset
+from diffusers import DDPMPipeline, UNet2DModel, DDPMScheduler
+from diffusers.optimization import get_cosine_schedule_with_warmup
+from diffusers.utils import make_image_grid
+from huggingface_hub import create_repo, upload_folder
+from peft import LoraConfig, get_peft_model
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+
+from config import TrainingConfig
+
+
+"""
+Or diffusion for simple images (cifar10 or fashion-mnist or mnist) and explore subtly different
+x_T's and what the output is. 
+
+Denoise each x_T multiple times to get a better picture of the distribution.
+Maybe use a set sequence of seeds for every denoising run (torch.Generator(seed=__)).
+
+Inter-concept space. Conciousness.
+"""
+
+
+def evaluate(config, epoch, pipeline):
+    # Sample some images from random noise (this is the backward diffusion process).
+    # The default pipeline output type is `List[PIL.Image]`
+    images = pipeline(
+        batch_size=config.eval_batch_size,
+        generator=torch.manual_seed(config.seed),
+        num_inference_steps=50
+    ).images
+
+    # Make a grid out of the images
+    image_grid = make_image_grid(images, rows=2, cols=2)
+
+    # Save the images
+    test_dir = Path(config.output_dir) / 'samples'
+    test_dir.mkdir(exist_ok=True)
+    image_grid.save(test_dir / f'{epoch:04d}.png')
+
+
+def print_trainable_parameters(model):
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    print(
+        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param:.2f}"
+    )
+
+
+if __name__ == '__main__':
+
+    config = TrainingConfig()
+    config.dataset_name = 'keremberke/painting-style-classification'
+
+    ds_dict = load_dataset(config.dataset_name, name='full')
+
+    preprocess = transforms.Compose([
+        transforms.Resize((config.image_size, config.image_size)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+    ])
+
+    def transform(examples):
+        return {
+            'images': [preprocess(img.convert('RGB')) for img in examples['image']]
+        }
+    
+    ds_dict.set_transform(transform)  # automatically applies preprocessing to samples as we load them
+
+    train_dataloader = torch.utils.data.DataLoader(ds_dict['train'], batch_size=config.train_batch_size, shuffle=True)
+    valid_dataloader = torch.utils.data.DataLoader(ds_dict['validation'], batch_size=config.eval_batch_size, shuffle=False)
+    test_dataloader = torch.utils.data.DataLoader(ds_dict['test'], batch_size=config.eval_batch_size, shuffle=False)
+
+    """pipe = DDPMPipeline.from_pretrained(
+        'google/ddpm-celebahq-256',
+        #use_safetensors=True
+    ).to('mps')
+    pipe.enable_attention_slicing()"""
+
+    unet = UNet2DModel.from_pretrained(
+        'google/ddpm-celebahq-256',
+        safetensors=True
+    ).to('mps')
+
+    scheduler = DDPMScheduler.from_pretrained(
+        'google/ddpm-celebahq-256'
+    )
+
+    lora_config = LoraConfig(r=8, lora_alpha=8, target_modules=['to_k','to_v'], lora_dropout=0.1, bias='none')
+    lora_unet = get_peft_model(unet, lora_config)
+
+    print_trainable_parameters(lora_unet)
+
+    optimizer = torch.optim.AdamW(lora_unet.parameters(), lr=config.learning_rate)
+    lr_scheduler = get_cosine_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=config.lr_warmup_steps,
+        num_training_steps=(len(train_dataloader) * config.num_epochs)
+    )
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=config.gradient_accumulation_steps,
+        mixed_precision=config.mixed_precision,
+        log_with='tensorboard',
+        project_dir=Path(config.output_dir) / 'logs'
+    )
+
+    if accelerator.is_main_process:
+        if config.push_to_hub:
+            repo_id = create_repo(repo_id=config.hub_model_id, exist_ok=True).repo_id
+
+        accelerator.init_trackers('ddpm-paintings-128-finetuned-cifar10')
+
+
+    epoch = -1
+    pipeline = DDPMPipeline(unet=accelerator.unwrap_model(lora_unet), scheduler=scheduler)
+    
+    upload_folder(
+        repo_id=repo_id,
+        folder_path=Path(config.output_dir).parent,
+        commit_message=f'Files: Epoch {epoch}',
+        ignore_patterns=['step_*', 'epoch_*'],
+        token='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'
+    )
+
+    pipeline.push_to_hub(
+        repo_id=config.hub_model_id,
+        commit_message=f'Model: Epoch {epoch}',
+        token='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'
+    )
+
+    exit()
+    global_step = 0
+    
+    for epoch in range(config.num_epochs):
+        pbar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)
+        pbar.set_description(f'Epoch {epoch}')
+
+        for idx, batch in enumerate(train_dataloader):
+            clean_images = batch['images'].to('mps')
+            
+            noise = torch.randn(clean_images.shape, device=clean_images.device)
+            bs = clean_images.shape[0]
+
+            ts = torch.randint(0, scheduler.config.num_train_timesteps, (bs,), device=clean_images.device, dtype=torch.int64)
+
+            noisy_images = scheduler.add_noise(clean_images, noise, ts)
+
+            with accelerator.accumulate(unet):
+                noise_pred = lora_unet(noisy_images, ts, return_dict=False)[0]
+                loss = F.mse_loss(noise_pred, noise)
+                accelerator.backward(loss)
+
+                accelerator.clip_grad_norm_(lora_unet.parameters(), 1.0)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            logs = {'loss': loss.detach().item(), 'lr': lr_scheduler.get_last_lr()[0], 'step': global_step}
+            pbar.update(1)
+            pbar.set_postfix(loss=logs['loss'], step=idx + 1)
+            accelerator.log(logs, step=global_step)
+            global_step += 1
+        
+        pbar.close()
+
+        if accelerator.is_main_process:
+            pipeline = DDPMPipeline(unet=accelerator.unwrap_model(lora_unet), scheduler=scheduler)
+
+            if (epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1:
+                # Save some images for model trained at end of epoch
+                evaluate(config, epoch, pipeline)
+
+            if (epoch + 1) % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:
+                if config.push_to_hub:
+                    upload_folder(
+                        repo_id=repo_id,
+                        folder_path=Path(config.output_dir).parent,
+                        commit_message=f'Files: Epoch {epoch}',
+                        ignore_patterns=['step_*', 'epoch_*'],
+                        token='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'
+                    )
+
+                    pipeline.push_to_hub(
+                        repo_id=config.hub_model_id,
+                        commit_message=f'Model: Epoch {epoch}',
+                        token='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'
+                    )
+                else:
+                    pipeline.save_pretrained(config.output_dir)

unet.txt

@@ -0,0 +1,319 @@
+UNet2DModel(
+  (conv_in): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+  (time_proj): Timesteps()
+  (time_embedding): TimestepEmbedding(
+    (linear_1): LoRACompatibleLinear(in_features=128, out_features=512, bias=True)
+    (act): SiLU()
+    (linear_2): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+  )
+  (down_blocks): ModuleList(
+    (0-1): 2 x DownBlock2D(
+      (resnets): ModuleList(
+        (0-1): 2 x ResnetBlock2D(
+          (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=128, bias=True)
+          (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+        )
+      )
+      (downsamplers): ModuleList(
+        (0): Downsample2D(
+          (conv): LoRACompatibleConv(128, 128, kernel_size=(3, 3), stride=(2, 2))
+        )
+      )
+    )
+    (2): DownBlock2D(
+      (resnets): ModuleList(
+        (0): ResnetBlock2D(
+          (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=256, bias=True)
+          (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(128, 256, kernel_size=(1, 1), stride=(1, 1))
+        )
+        (1): ResnetBlock2D(
+          (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=256, bias=True)
+          (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+        )
+      )
+      (downsamplers): ModuleList(
+        (0): Downsample2D(
+          (conv): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(2, 2))
+        )
+      )
+    )
+    (3): DownBlock2D(
+      (resnets): ModuleList(
+        (0-1): 2 x ResnetBlock2D(
+          (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=256, bias=True)
+          (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+        )
+      )
+      (downsamplers): ModuleList(
+        (0): Downsample2D(
+          (conv): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(2, 2))
+        )
+      )
+    )
+    (4): AttnDownBlock2D(
+      (attentions): ModuleList(
+        (0-1): 2 x Attention(
+          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (to_q): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (to_k): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (to_v): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (to_out): ModuleList(
+            (0): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+            (1): Dropout(p=0.0, inplace=False)
+          )
+        )
+      )
+      (resnets): ModuleList(
+        (0): ResnetBlock2D(
+          (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(256, 512, kernel_size=(1, 1), stride=(1, 1))
+        )
+        (1): ResnetBlock2D(
+          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+        )
+      )
+      (downsamplers): ModuleList(
+        (0): Downsample2D(
+          (conv): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(2, 2))
+        )
+      )
+    )
+    (5): DownBlock2D(
+      (resnets): ModuleList(
+        (0-1): 2 x ResnetBlock2D(
+          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+        )
+      )
+    )
+  )
+  (up_blocks): ModuleList(
+    (0): UpBlock2D(
+      (resnets): ModuleList(
+        (0-2): 3 x ResnetBlock2D(
+          (norm1): GroupNorm(32, 1024, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(1024, 512, kernel_size=(1, 1), stride=(1, 1))
+        )
+      )
+      (upsamplers): ModuleList(
+        (0): Upsample2D(
+          (conv): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        )
+      )
+    )
+    (1): AttnUpBlock2D(
+      (attentions): ModuleList(
+        (0-2): 3 x Attention(
+          (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (to_q): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (to_k): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (to_v): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (to_out): ModuleList(
+            (0): LoRACompatibleLinear(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, 1024, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(1024, 512, kernel_size=(1, 1), stride=(1, 1))
+        )
+        (2): ResnetBlock2D(
+          (norm1): GroupNorm(32, 768, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(768, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(768, 512, kernel_size=(1, 1), stride=(1, 1))
+        )
+      )
+      (upsamplers): ModuleList(
+        (0): Upsample2D(
+          (conv): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        )
+      )
+    )
+    (2): UpBlock2D(
+      (resnets): ModuleList(
+        (0): ResnetBlock2D(
+          (norm1): GroupNorm(32, 768, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(768, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=256, bias=True)
+          (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(768, 256, kernel_size=(1, 1), stride=(1, 1))
+        )
+        (1-2): 2 x ResnetBlock2D(
+          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=256, bias=True)
+          (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(512, 256, kernel_size=(1, 1), stride=(1, 1))
+        )
+      )
+      (upsamplers): ModuleList(
+        (0): Upsample2D(
+          (conv): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        )
+      )
+    )
+    (3): UpBlock2D(
+      (resnets): ModuleList(
+        (0-1): 2 x ResnetBlock2D(
+          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=256, bias=True)
+          (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(512, 256, kernel_size=(1, 1), stride=(1, 1))
+        )
+        (2): ResnetBlock2D(
+          (norm1): GroupNorm(32, 384, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=256, bias=True)
+          (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(384, 256, kernel_size=(1, 1), stride=(1, 1))
+        )
+      )
+      (upsamplers): ModuleList(
+        (0): Upsample2D(
+          (conv): LoRACompatibleConv(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        )
+      )
+    )
+    (4): UpBlock2D(
+      (resnets): ModuleList(
+        (0): ResnetBlock2D(
+          (norm1): GroupNorm(32, 384, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(384, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=128, bias=True)
+          (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(384, 128, kernel_size=(1, 1), stride=(1, 1))
+        )
+        (1-2): 2 x ResnetBlock2D(
+          (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=128, bias=True)
+          (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(256, 128, kernel_size=(1, 1), stride=(1, 1))
+        )
+      )
+      (upsamplers): ModuleList(
+        (0): Upsample2D(
+          (conv): LoRACompatibleConv(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        )
+      )
+    )
+    (5): UpBlock2D(
+      (resnets): ModuleList(
+        (0-2): 3 x ResnetBlock2D(
+          (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)
+          (conv1): LoRACompatibleConv(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=128, bias=True)
+          (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)
+          (dropout): Dropout(p=0.0, inplace=False)
+          (conv2): LoRACompatibleConv(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+          (nonlinearity): SiLU()
+          (conv_shortcut): LoRACompatibleConv(256, 128, kernel_size=(1, 1), stride=(1, 1))
+        )
+      )
+    )
+  )
+  (mid_block): UNetMidBlock2D(
+    (attentions): ModuleList(
+      (0): Attention(
+        (group_norm): GroupNorm(32, 512, eps=1e-06, affine=True)
+        (to_q): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+        (to_k): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+        (to_v): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+        (to_out): ModuleList(
+          (0): LoRACompatibleLinear(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): LoRACompatibleConv(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        (time_emb_proj): LoRACompatibleLinear(in_features=512, out_features=512, bias=True)
+        (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)
+        (dropout): Dropout(p=0.0, inplace=False)
+        (conv2): LoRACompatibleConv(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))
+)