Upload DDPMPipeline

config.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from pathlib import Path
 
 
 @dataclass
@@ -13,10 +14,10 @@ class TrainingConfig:
     save_image_epochs = 1
     save_model_epochs = 3
     mixed_precision = 'fp16'  # `no` for float32, `fp16` for automatic mixed precision
-    output_dir = '.'
+    output_dir = str(Path(__file__).parent)
 
     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_model_id = 'jmemon/ddpm-paintings-128-finetuned-celebahq'  # 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

logs/ddpm-paintings-128-finetuned-celebahq/events.out.tfevents.1701732706.coffee.29296.0

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

logs/ddpm-paintings-128-finetuned-celebahq/events.out.tfevents.1701734614.coffee.30162.0

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

logs/ddpm-paintings-128-finetuned-celebahq/events.out.tfevents.1701734735.coffee.30262.0

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

logs/ddpm-paintings-128-finetuned-celebahq/events.out.tfevents.1701735986.coffee.30996.0

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

logs/ddpm-paintings-128-finetuned-celebahq/events.out.tfevents.1701744188.coffee.33147.0

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

logs/ddpm-paintings-128-finetuned-celebahq/events.out.tfevents.1701745393.coffee.33519.0

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

logs/ddpm-paintings-128-finetuned-cifar10/events.out.tfevents.1701726567.coffee.25846.0

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

logs/ddpm-paintings-128-finetuned-cifar10/events.out.tfevents.1701726710.coffee.25924.0

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

logs/ddpm-paintings-128-finetuned-cifar10/events.out.tfevents.1701728614.coffee.26917.0

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

logs/ddpm-paintings-128-finetuned-cifar10/events.out.tfevents.1701729930.coffee.27724.0

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

logs/ddpm-paintings-128-finetuned-cifar10/events.out.tfevents.1701731364.coffee.28580.0

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

main.py

@@ -17,7 +17,7 @@ from config import TrainingConfig
 
 
 """
-Or diffusion for simple images (cifar10 or fashion-mnist or mnist) and explore subtly different
+Or diffusion for simple images 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.
@@ -81,27 +81,37 @@ if __name__ == '__main__':
     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)
 
+    """
     unet = UNet2DModel.from_pretrained(
-        'google/ddpm-celebahq-256',
-        safetensors=True
+        'google/ddpm-celebahq-256'
     ).to('mps')
     
     scheduler = DDPMScheduler.from_pretrained(
         'google/ddpm-celebahq-256'
     )
+    """
 
-    """unet=UNet2DModel.from_pretrained(
+    """
+    unet=UNet2DModel.from_pretrained(
         'jmemon/ddpm-paintings-128-finetuned-celebahq'
     ).to('mps')
 
     scheduler = DDPMScheduler.from_pretrained(
         'jmemon/ddpm-paintings-128-finetuned-celebahq'
-    )"""
+    )
+    """
+
+    unet = UNet2DModel.from_pretrained(
+        str(Path(__file__).parent / 'unet')
+    ).to('mps')
+
+    scheduler = DDPMScheduler.from_pretrained(
+        str(Path(__file__).parent / 'scheduler')
+    )
 
     lora_config = LoraConfig(
         r=8, 
         lora_alpha=8, 
-        #modules_to_save=['model'], 
         target_modules=['to_k','to_v'], 
         lora_dropout=0.1, 
         bias='none')
@@ -127,11 +137,11 @@ if __name__ == '__main__':
         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')
+        accelerator.init_trackers('ddpm-paintings-128-finetuned-celebahq')
 
     global_step = 0
     
-    for epoch in range(config.num_epochs):
+    for epoch in range(3, config.num_epochs + 3):
         pbar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)
         pbar.set_description(f'Epoch {epoch}')
 
@@ -145,7 +155,7 @@ if __name__ == '__main__':
 
             noisy_images = scheduler.add_noise(clean_images, noise, ts)
 
-            with accelerator.accumulate(unet):
+            with accelerator.accumulate(lora_unet):
                 noise_pred = lora_unet(noisy_images, ts, return_dict=False)[0]
                 loss = F.mse_loss(noise_pred, noise)
                 accelerator.backward(loss)
@@ -164,6 +174,7 @@ if __name__ == '__main__':
         pbar.close()
 
         if accelerator.is_main_process:
+            #pipeline = DDPMPipeline(unet=accelerator.unwrap_model(lora_unet).merge_and_unload(), scheduler=scheduler)
             pipeline = DDPMPipeline(unet=accelerator.unwrap_model(lora_unet), scheduler=scheduler)
 
             if (epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1:
@@ -171,16 +182,29 @@ if __name__ == '__main__':
                 evaluate(config, epoch, pipeline)
 
             if (epoch + 1) % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:
+                _pipeline = DDPMPipeline(
+                    unet=accelerator.unwrap_model(lora_unet).merge_and_unload(), 
+                    scheduler=scheduler)
                 if config.push_to_hub:
-                    torch.save(pipeline.unet.state_dict(), Path(config.output_dir) / 'model' / 'adapter_model.bin')
+                    _pipeline.save_pretrained(
+                        config.output_dir, 
+                        push_to_hub=True, 
+                        repo_id=repo_id, 
+                        token='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'
+                    )
 
                     upload_folder(
                         repo_id=repo_id,
-                        folder_path=Path(config.output_dir).parent,
+                        folder_path=config.output_dir,
                         commit_message=f'Epoch {epoch}',
-                        ignore_patterns=['logs', '.DS_Store'],
+                        ignore_patterns=['logs/*', '*/.DS_Store'],
                         token='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'
                     )
 
+                    model_loc = 'jmemon/ddpm-paintings-128-finetuned-celebahq'
                 else:
-                    pipeline.save_pretrained(config.output_dir)
+                    _pipeline.save_pretrained(config.output_dir)
+                    model_loc = str(Path(__file__).parent / 'diffusion_model_pytorch.bin')
+
+                unet = UNet2DModel.from_pretrained(model_loc)
+                lora_unet = get_peft_model(unet, lora_config)

scheduler/scheduler_config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "DDPMScheduler",
+  "_diffusers_version": "0.24.0",
+  "beta_end": 0.02,
+  "beta_schedule": "linear",
+  "beta_start": 0.0001,
+  "clip_sample": true,
+  "clip_sample_range": 1.0,
+  "dynamic_thresholding_ratio": 0.995,
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "sample_max_value": 1.0,
+  "steps_offset": 0,
+  "thresholding": false,
+  "timestep_spacing": "leading",
+  "trained_betas": null,
+  "variance_type": "fixed_small"
+}

unet/config.json

@@ -0,0 +1,52 @@
+{
+  "_class_name": "UNet2DModel",
+  "_diffusers_version": "0.24.0",
+  "_name_or_path": "/Users/jpmemon/Desktop/School/Fall-2023/Capstone/DiffusionEx/unet",
+  "act_fn": "silu",
+  "add_attention": true,
+  "attention_head_dim": null,
+  "attn_norm_num_groups": null,
+  "block_out_channels": [
+    128,
+    128,
+    256,
+    256,
+    512,
+    512
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "down_block_types": [
+    "DownBlock2D",
+    "DownBlock2D",
+    "DownBlock2D",
+    "DownBlock2D",
+    "AttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 0,
+  "downsample_type": "conv",
+  "dropout": 0.0,
+  "flip_sin_to_cos": false,
+  "freq_shift": 1,
+  "in_channels": 3,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "norm_eps": 1e-06,
+  "norm_num_groups": 32,
+  "num_class_embeds": null,
+  "num_train_timesteps": null,
+  "out_channels": 3,
+  "resnet_time_scale_shift": "default",
+  "sample_size": 256,
+  "time_embedding_type": "positional",
+  "up_block_types": [
+    "UpBlock2D",
+    "AttnUpBlock2D",
+    "UpBlock2D",
+    "UpBlock2D",
+    "UpBlock2D",
+    "UpBlock2D"
+  ],
+  "upsample_type": "conv"
+}

unet/diffusion_pytorch_model.safetensors

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