{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# DIFFUSION","metadata":{}},{"cell_type":"code","source":"mkdir -p /kaggle/working/data","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:11.702215Z","iopub.execute_input":"2023-05-27T18:18:11.703003Z","iopub.status.idle":"2023-05-27T18:18:12.821691Z","shell.execute_reply.started":"2023-05-27T18:18:11.702962Z","shell.execute_reply":"2023-05-27T18:18:12.820146Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cp -r /kaggle/input/landscape-recognition-image-dataset-12k-images/'Landscape Classification'/'Landscape Classification'/'Training Data'/Coast /kaggle/working/data ","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:12.824098Z","iopub.execute_input":"2023-05-27T18:18:12.831046Z","iopub.status.idle":"2023-05-27T18:18:15.634505Z","shell.execute_reply.started":"2023-05-27T18:18:12.830993Z","shell.execute_reply":"2023-05-27T18:18:15.633093Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"custom data for now, will change?","metadata":{}},{"cell_type":"code","source":"from dataclasses import dataclass\nimport torch\nimport os\n\ndef get_default_device():\n    \"\"\"Pick GPU if available, else CPU\"\"\"\n    if torch.cuda.is_available():\n        return torch.device('cuda')\n    else:\n        return torch.device('cpu')\n \n@dataclass\nclass BaseConfig:\n    DATASET = \"Monet\" # \"Flowers\", \"MNIST\", \"Cifar-10\", \"Flowers\"\n    DEVICE = get_default_device()\n    # For logging inferece images and saving checkpoints.\n    root_log_dir = os.path.join(\"Logs_Checkpoints\", \"Inference\")\n    root_checkpoint_dir = os.path.join(\"Logs_Checkpoints\", \"checkpoints\")\n \n    # Current log and checkpoint directory.\n    log_dir = \"version_0\"\n    checkpoint_dir = \"version_0\"\n \n \n@dataclass\nclass TrainingConfig:\n    TIMESTEPS = 800  # Define number of diffusion timesteps\n    IMG_SHAPE = (1, 32, 32) if BaseConfig.DATASET == \"MNIST\" else (3, 256, 256)\n    NUM_EPOCHS = 500\n    BATCH_SIZE = 32\n    LR = 2e-4\n    NUM_WORKERS = 2\n    \n    \ndef to_device(data, device):\n    \"\"\"Move tensor(s) to chosen device\"\"\"\n    if isinstance(data, (list,tuple)):\n        return [to_device(x, device) for x in data]\n    return data.to(device, non_blocking=True)\n\n\nclass DeviceDataLoader():\n    \"\"\"Wrap a dataloader to move data to a device\"\"\"\n    def __init__(self, dl, device):\n        self.dl = dl\n        self.device = device\n        \n    def __iter__(self):\n        \"\"\"Yield a batch of data after moving it to device\"\"\"\n        for b in self.dl: \n            yield to_device(b, self.device)\n\n    def __len__(self):\n        \"\"\"Number of batches\"\"\"\n        return len(self.dl)","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:15.636772Z","iopub.execute_input":"2023-05-27T18:18:15.637179Z","iopub.status.idle":"2023-05-27T18:18:15.654620Z","shell.execute_reply.started":"2023-05-27T18:18:15.637135Z","shell.execute_reply":"2023-05-27T18:18:15.653552Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torchvision\nimport torchvision.transforms as TF\nimport torchvision.datasets as datasets\nfrom torch.utils.data import Dataset, DataLoader\n \n \ndef get_dataset(dataset_name='MNIST'):\n    transforms = torchvision.transforms.Compose(\n        [\n            torchvision.transforms.ToTensor(),\n            torchvision.transforms.Resize((64, 64), \n                                          interpolation=torchvision.transforms.InterpolationMode.BICUBIC, \n                                          antialias=True),\n            torchvision.transforms.RandomHorizontalFlip(),\n#             torchvision.transforms.Normalize(MEAN, STD),\n            torchvision.transforms.Lambda(lambda t: (t * 2) - 1) # Scale between [-1, 1] \n        ]\n    )\n     \n    if dataset_name.upper() == \"MNIST\":\n        return datasets.MNIST(root=\"data\", train=True, download=True, transform=transforms)\n    elif dataset_name == \"Cifar-10\":    \n        return datasets.CIFAR10(root=\"data\", train=True, download=True, transform=transforms)\n    elif dataset_name == \"Cifar-100\":\n        return datasets.CIFAR10(root=\"data\", train=True, download=True, transform=transforms)\n    elif dataset_name == \"Flowers\":\n        return datasets.ImageFolder(root=\"/kaggle/input/flowers-recognition/flowers\", transform=transforms)\n    elif dataset_name == \"Monet\":\n        # path here - dirs have to be nested\n        print(\"we are here\")\n        return datasets.ImageFolder(root=\"/kaggle/working/data\", transform=transforms)\n \ndef inverse_transform(tensors):\n    \"\"\"Convert tensors from [-1., 1.] to [0., 255.]\"\"\"\n    return ((tensors.clamp(-1, 1) + 1.0) / 2.0) * 255.0","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:15.658435Z","iopub.execute_input":"2023-05-27T18:18:15.658870Z","iopub.status.idle":"2023-05-27T18:18:15.671688Z","shell.execute_reply.started":"2023-05-27T18:18:15.658830Z","shell.execute_reply":"2023-05-27T18:18:15.670489Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"DATA LOADER:","metadata":{}},{"cell_type":"code","source":"def get_dataloader(dataset_name='MNIST', \n                   batch_size=32, \n                   pin_memory=False, \n                   shuffle=True, \n                   num_workers=0, \n                   device=\"cpu\"\n                  ):\n    dataset      = get_dataset(dataset_name=dataset_name)\n    dataloader = DataLoader(dataset, batch_size=batch_size, \n                            pin_memory=pin_memory, \n                            num_workers=num_workers, \n                            shuffle=shuffle\n                           )\n    # Used for moving batch of data to the user-specified machine: cpu or gpu\n    device_dataloader = DeviceDataLoader(dataloader, device)\n    return device_dataloader","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:15.673455Z","iopub.execute_input":"2023-05-27T18:18:15.673901Z","iopub.status.idle":"2023-05-27T18:18:15.685299Z","shell.execute_reply.started":"2023-05-27T18:18:15.673862Z","shell.execute_reply":"2023-05-27T18:18:15.684308Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loader = get_dataloader(\n    dataset_name=BaseConfig.DATASET,\n    batch_size=32,\n    device=\"cpu\",\n)","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:15.686839Z","iopub.execute_input":"2023-05-27T18:18:15.687435Z","iopub.status.idle":"2023-05-27T18:18:15.712962Z","shell.execute_reply.started":"2023-05-27T18:18:15.687391Z","shell.execute_reply":"2023-05-27T18:18:15.711871Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# visualize dataset","metadata":{}},{"cell_type":"code","source":"from torchvision.utils import make_grid\nimport matplotlib.pyplot as plt\n \nplt.figure(figsize=(10, 4), facecolor='white')\n \nfor b_image, _ in loader:\n    b_image = inverse_transform(b_image)\n    grid_img = make_grid(b_image / 255.0, nrow=16, padding=True, pad_value=1)\n    plt.imshow(grid_img.permute(1, 2, 0))\n    plt.axis(\"off\")\n    break","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:15.714554Z","iopub.execute_input":"2023-05-27T18:18:15.714905Z","iopub.status.idle":"2023-05-27T18:18:15.974969Z","shell.execute_reply.started":"2023-05-27T18:18:15.714868Z","shell.execute_reply":"2023-05-27T18:18:15.973861Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"helepr","metadata":{}},{"cell_type":"code","source":"def get(element: torch.Tensor, t: torch.Tensor):\n    \"\"\"\n    Get value at index position \"t\" in \"element\" and\n        reshape it to have the same dimension as a batch of images.\n    \"\"\"\n    ele = element.gather(-1, t)\n    return ele.reshape(-1, 1, 1, 1)","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:15.976520Z","iopub.execute_input":"2023-05-27T18:18:15.977594Z","iopub.status.idle":"2023-05-27T18:18:15.984578Z","shell.execute_reply.started":"2023-05-27T18:18:15.977552Z","shell.execute_reply":"2023-05-27T18:18:15.982923Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Diffusing images","metadata":{}},{"cell_type":"code","source":"class SimpleDiffusion:\n    def __init__(\n        self,\n        num_diffusion_timesteps=1000,\n        img_shape=(3, 256, 256),\n        device=\"cpu\",\n    ):\n        self.num_diffusion_timesteps = num_diffusion_timesteps\n        self.img_shape = img_shape\n        self.device = device\n        self.initialize()\n \n    def initialize(self):\n        # BETAs & ALPHAs required at different places in the Algorithm.\n        self.beta  = self.get_betas()\n        self.alpha = 1 - self.beta\n         \n        self_sqrt_beta                       = torch.sqrt(self.beta)\n        self.alpha_cumulative                = torch.cumprod(self.alpha, dim=0)\n        self.sqrt_alpha_cumulative           = torch.sqrt(self.alpha_cumulative)\n        self.one_by_sqrt_alpha               = 1. / torch.sqrt(self.alpha)\n        self.sqrt_one_minus_alpha_cumulative = torch.sqrt(1 - self.alpha_cumulative)\n          \n    def get_betas(self):\n        \"\"\"linear schedule, proposed in original ddpm paper\"\"\"\n        scale = 1000 / self.num_diffusion_timesteps\n        beta_start = scale * 1e-4\n        beta_end = scale * 0.02\n        return torch.linspace(\n            beta_start,\n            beta_end,\n            self.num_diffusion_timesteps,\n            dtype=torch.float32,\n            device=self.device,\n        )","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:15.986832Z","iopub.execute_input":"2023-05-27T18:18:15.988677Z","iopub.status.idle":"2023-05-27T18:18:16.002507Z","shell.execute_reply.started":"2023-05-27T18:18:15.988543Z","shell.execute_reply":"2023-05-27T18:18:16.001272Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# creation of forward diffusion","metadata":{}},{"cell_type":"markdown","source":"# now visuals","metadata":{}},{"cell_type":"code","source":"def forward_diffusion(sd: SimpleDiffusion, x0: torch.Tensor, timesteps: torch.Tensor):\n    eps = torch.randn_like(x0)  # Noise\n    mean    = get(sd.sqrt_alpha_cumulative, t=timesteps) * x0  # Image scaled\n    std_dev = get(sd.sqrt_one_minus_alpha_cumulative, t=timesteps) # Noise scaled\n    sample  = mean + std_dev * eps # scaled inputs * scaled noise\n \n    return sample, eps \n\nsd = SimpleDiffusion(num_diffusion_timesteps=TrainingConfig.TIMESTEPS, device=\"cpu\")\n \nloader = iter( # to iterate\n    get_dataloader(\n        dataset_name=BaseConfig.DATASET,\n        batch_size=6,\n        device=\"cpu\",\n    )\n)","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:16.008279Z","iopub.execute_input":"2023-05-27T18:18:16.009320Z","iopub.status.idle":"2023-05-27T18:18:16.033233Z","shell.execute_reply.started":"2023-05-27T18:18:16.009279Z","shell.execute_reply":"2023-05-27T18:18:16.032056Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"define timesteps","metadata":{}},{"cell_type":"code","source":"x0s, _ = next(loader)\n \nnoisy_images = []\n#specific_timesteps = [0, 10, 50, 100, 150, 200, 250, 300, 400, 600, 800, 999]\nspecific_timesteps = [0, 10, 50, 100, 150, 200, 250, 300, 400, 600, 799]\n    \nfor timestep in specific_timesteps:\n    timestep = torch.as_tensor(timestep, dtype=torch.long)\n \n    xts, _ = forward_diffusion(sd, x0s, timestep)\n    xts    = inverse_transform(xts) / 255.0\n    xts    = make_grid(xts, nrow=1, padding=1)\n     \n    noisy_images.append(xts)","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:16.034756Z","iopub.execute_input":"2023-05-27T18:18:16.035439Z","iopub.status.idle":"2023-05-27T18:18:16.071582Z","shell.execute_reply.started":"2023-05-27T18:18:16.035393Z","shell.execute_reply":"2023-05-27T18:18:16.070586Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"hurra?","metadata":{}},{"cell_type":"code","source":"_, ax = plt.subplots(1, len(noisy_images), figsize=(10, 5), facecolor='white')\n \nfor i, (timestep, noisy_sample) in enumerate(zip(specific_timesteps, noisy_images)):\n    ax[i].imshow(noisy_sample.squeeze(0).permute(1, 2, 0))\n    ax[i].set_title(f\"t={timestep}\", fontsize=8)\n    ax[i].axis(\"off\")\n    ax[i].grid(False)\n \nplt.suptitle(\"Forward Diffusion Process\", y=0.9)\nplt.axis(\"off\")\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:16.072884Z","iopub.execute_input":"2023-05-27T18:18:16.073571Z","iopub.status.idle":"2023-05-27T18:18:16.848794Z","shell.execute_reply.started":"2023-05-27T18:18:16.073531Z","shell.execute_reply":"2023-05-27T18:18:16.847066Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# simple train","metadata":{}},{"cell_type":"markdown","source":"UNET","metadata":{}},{"cell_type":"code","source":"import torch.nn as nn\n\nclass SinusoidalPositionEmbeddings(nn.Module):\n    def __init__(self, total_time_steps=1000, time_emb_dims=128, time_emb_dims_exp=512):\n        super().__init__()\n\n        half_dim = time_emb_dims // 2\n\n        emb = math.log(10000) / (half_dim - 1)\n        emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)\n\n        ts = torch.arange(total_time_steps, dtype=torch.float32)\n\n        emb = torch.unsqueeze(ts, dim=-1) * torch.unsqueeze(emb, dim=0)\n        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)\n\n        self.time_blocks = nn.Sequential(\n            nn.Embedding.from_pretrained(emb),\n            nn.Linear(in_features=time_emb_dims, out_features=time_emb_dims_exp),\n            nn.SiLU(),\n            nn.Linear(in_features=time_emb_dims_exp, out_features=time_emb_dims_exp),\n        )\n\n    def forward(self, time):\n        return self.time_blocks(time)\n\n\nclass AttentionBlock(nn.Module):\n    def __init__(self, channels=64):\n        super().__init__()\n        self.channels = channels\n\n        self.group_norm = nn.GroupNorm(num_groups=8, num_channels=channels)\n        self.mhsa = nn.MultiheadAttention(embed_dim=self.channels, num_heads=4, batch_first=True)\n\n    def forward(self, x):\n        B, _, H, W = x.shape\n        h = self.group_norm(x)\n        h = h.reshape(B, self.channels, H * W).swapaxes(1, 2)  # [B, C, H, W] --> [B, C, H * W] --> [B, H*W, C]\n        h, _ = self.mhsa(h, h, h)  # [B, H*W, C]\n        h = h.swapaxes(2, 1).view(B, self.channels, H, W)  # [B, C, H*W] --> [B, C, H, W]\n        return x + h\n\n\nclass ResnetBlock(nn.Module):\n    def __init__(self, *, in_channels, out_channels, dropout_rate=0.1, time_emb_dims=512, apply_attention=False):\n        super().__init__()\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n\n        self.act_fn = nn.SiLU()\n        # Group 1\n        self.normlize_1 = nn.GroupNorm(num_groups=8, num_channels=self.in_channels)\n        self.conv_1 = nn.Conv2d(in_channels=self.in_channels, out_channels=self.out_channels, kernel_size=3, stride=1, padding=\"same\")\n\n        # Group 2 time embedding\n        self.dense_1 = nn.Linear(in_features=time_emb_dims, out_features=self.out_channels)\n\n        # Group 3\n        self.normlize_2 = nn.GroupNorm(num_groups=8, num_channels=self.out_channels)\n        self.dropout = nn.Dropout2d(p=dropout_rate)\n        self.conv_2 = nn.Conv2d(in_channels=self.out_channels, out_channels=self.out_channels, kernel_size=3, stride=1, padding=\"same\")\n\n        if self.in_channels != self.out_channels:\n            self.match_input = nn.Conv2d(in_channels=self.in_channels, out_channels=self.out_channels, kernel_size=1, stride=1)\n        else:\n            self.match_input = nn.Identity()\n\n        if apply_attention:\n            self.attention = AttentionBlock(channels=self.out_channels)\n        else:\n            self.attention = nn.Identity()\n\n    def forward(self, x, t):\n        # group 1\n        h = self.act_fn(self.normlize_1(x))\n        h = self.conv_1(h)\n\n        # group 2\n        # add in timestep embedding\n        h += self.dense_1(self.act_fn(t))[:, :, None, None]\n\n        # group 3\n        h = self.act_fn(self.normlize_2(h))\n        h = self.dropout(h)\n        h = self.conv_2(h)\n\n        # Residual and attention\n        h = h + self.match_input(x)\n        h = self.attention(h)\n\n        return h\n\n\nclass DownSample(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.downsample = nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=3, stride=2, padding=1)\n\n    def forward(self, x, *args):\n        return self.downsample(x)\n\n\nclass UpSample(nn.Module):\n    def __init__(self, in_channels):\n        super().__init__()\n\n        self.upsample = nn.Sequential(\n            nn.Upsample(scale_factor=2, mode=\"nearest\"),\n            nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=3, stride=1, padding=1),\n        )\n\n    def forward(self, x, *args):\n        return self.upsample(x)\n\n\nclass UNet(nn.Module):\n    def __init__(\n        self,\n        input_channels=3,\n        output_channels=3,\n        num_res_blocks=2,\n        base_channels=128,\n        base_channels_multiples=(1, 2, 4, 8),\n        apply_attention=(False, False, True, False),\n        dropout_rate=0.1,\n        time_multiple=4,\n    ):\n        super().__init__()\n\n        time_emb_dims_exp = base_channels * time_multiple\n        self.time_embeddings = SinusoidalPositionEmbeddings(time_emb_dims=base_channels, time_emb_dims_exp=time_emb_dims_exp)\n\n        self.first = nn.Conv2d(in_channels=input_channels, out_channels=base_channels, kernel_size=3, stride=1, padding=\"same\")\n\n        num_resolutions = len(base_channels_multiples)\n\n        # Encoder part of the UNet. Dimension reduction.\n        self.encoder_blocks = nn.ModuleList()\n        curr_channels = [base_channels]\n        in_channels = base_channels\n\n        for level in range(num_resolutions):\n            out_channels = base_channels * base_channels_multiples[level]\n\n            for _ in range(num_res_blocks):\n\n                block = ResnetBlock(\n                    in_channels=in_channels,\n                    out_channels=out_channels,\n                    dropout_rate=dropout_rate,\n                    time_emb_dims=time_emb_dims_exp,\n                    apply_attention=apply_attention[level],\n                )\n                self.encoder_blocks.append(block)\n\n                in_channels = out_channels\n                curr_channels.append(in_channels)\n\n            if level != (num_resolutions - 1):\n                self.encoder_blocks.append(DownSample(channels=in_channels))\n                curr_channels.append(in_channels)\n\n        # Bottleneck in between\n        self.bottleneck_blocks = nn.ModuleList(\n            (\n                ResnetBlock(\n                    in_channels=in_channels,\n                    out_channels=in_channels,\n                    dropout_rate=dropout_rate,\n                    time_emb_dims=time_emb_dims_exp,\n                    apply_attention=True,\n                ),\n                ResnetBlock(\n                    in_channels=in_channels,\n                    out_channels=in_channels,\n                    dropout_rate=dropout_rate,\n                    time_emb_dims=time_emb_dims_exp,\n                    apply_attention=False,\n                ),\n            )\n        )\n\n        # Decoder part of the UNet. Dimension restoration with skip-connections.\n        self.decoder_blocks = nn.ModuleList()\n\n        for level in reversed(range(num_resolutions)):\n            out_channels = base_channels * base_channels_multiples[level]\n\n            for _ in range(num_res_blocks + 1):\n                encoder_in_channels = curr_channels.pop()\n                block = ResnetBlock(\n                    in_channels=encoder_in_channels + in_channels,\n                    out_channels=out_channels,\n                    dropout_rate=dropout_rate,\n                    time_emb_dims=time_emb_dims_exp,\n                    apply_attention=apply_attention[level],\n                )\n\n                in_channels = out_channels\n                self.decoder_blocks.append(block)\n\n            if level != 0:\n                self.decoder_blocks.append(UpSample(in_channels))\n\n        self.final = nn.Sequential(\n            nn.GroupNorm(num_groups=8, num_channels=in_channels),\n            nn.SiLU(),\n            nn.Conv2d(in_channels=in_channels, out_channels=output_channels, kernel_size=3, stride=1, padding=\"same\"),\n        )\n\n    def forward(self, x, t):\n\n        time_emb = self.time_embeddings(t)\n\n        h = self.first(x)\n        outs = [h]\n\n        for layer in self.encoder_blocks:\n            h = layer(h, time_emb)\n            outs.append(h)\n\n        for layer in self.bottleneck_blocks:\n            h = layer(h, time_emb)\n\n        for layer in self.decoder_blocks:\n            if isinstance(layer, ResnetBlock):\n                out = outs.pop()\n                h = torch.cat([h, out], dim=1)\n            h = layer(h, time_emb)\n\n        h = self.final(h)\n\n        return h","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:16.850560Z","iopub.execute_input":"2023-05-27T18:18:16.851288Z","iopub.status.idle":"2023-05-27T18:18:16.904785Z","shell.execute_reply.started":"2023-05-27T18:18:16.851245Z","shell.execute_reply":"2023-05-27T18:18:16.903495Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"def train_one_epoch(model, sd, loader, optimizer, scaler, loss_fn, epoch=800, \n                   base_config=BaseConfig(), training_config=TrainingConfig()):\n     \n    loss_record = MeanMetric()\n    model.train()\n \n    with tqdm(total=len(loader), dynamic_ncols=True) as tq:\n        tq.set_description(f\"Train :: Epoch: {epoch}/{training_config.NUM_EPOCHS}\")\n          \n        for x0s, _ in loader: # line 1, 2\n            tq.update(1)\n             \n            ts = torch.randint(low=1, high=training_config.TIMESTEPS, size=(x0s.shape[0],), device=base_config.DEVICE) # line 3\n            xts, gt_noise = forward_diffusion(sd, x0s, ts) # line 4\n \n            with amp.autocast():\n                pred_noise = model(xts, ts)\n                loss = loss_fn(gt_noise, pred_noise) # line 5\n \n            optimizer.zero_grad(set_to_none=True)\n            scaler.scale(loss).backward()\n \n            # scaler.unscale_(optimizer)\n            # torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n \n            scaler.step(optimizer)\n            scaler.update()\n \n            loss_value = loss.detach().item()\n            loss_record.update(loss_value)\n \n            tq.set_postfix_str(s=f\"Loss: {loss_value:.4f}\")\n \n        mean_loss = loss_record.compute().item()\n     \n        tq.set_postfix_str(s=f\"Epoch Loss: {mean_loss:.4f}\")\n     \n    return mean_loss","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:16.906898Z","iopub.execute_input":"2023-05-27T18:18:16.907359Z","iopub.status.idle":"2023-05-27T18:18:16.921379Z","shell.execute_reply.started":"2023-05-27T18:18:16.907318Z","shell.execute_reply":"2023-05-27T18:18:16.920159Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Algorithm 2: Sampling\n     \n@torch.no_grad()\ndef reverse_diffusion(model, sd, timesteps=1000, img_shape=(3, 64, 64), \n                      num_images=5, nrow=8, device=\"cpu\", **kwargs):\n \n    x = torch.randn((num_images, *img_shape), device=device)\n    model.eval()\n \n    if kwargs.get(\"generate_video\", False):\n        outs = []\n \n    for time_step in tqdm(iterable=reversed(range(1, timesteps)), \n                          total=timesteps-1, dynamic_ncols=False, \n                          desc=\"Sampling :: \", position=0):\n \n        ts = torch.ones(num_images, dtype=torch.long, device=device) * time_step\n        z = torch.randn_like(x) if time_step > 1 else torch.zeros_like(x)\n \n        predicted_noise = model(x, ts)\n \n        beta_t                            = get(sd.beta, ts)\n        one_by_sqrt_alpha_t               = get(sd.one_by_sqrt_alpha, ts)\n        sqrt_one_minus_alpha_cumulative_t = get(sd.sqrt_one_minus_alpha_cumulative, ts) \n \n        x = (\n            one_by_sqrt_alpha_t\n            * (x - (beta_t / sqrt_one_minus_alpha_cumulative_t) * predicted_noise)\n            + torch.sqrt(beta_t) * z\n        )\n \n        if kwargs.get(\"generate_video\", False):\n            x_inv = inverse_transform(x).type(torch.uint8)\n            grid = make_grid(x_inv, nrow=nrow, pad_value=255.0).to(\"cpu\")\n            ndarr = torch.permute(grid, (1, 2, 0)).numpy()[:, :, ::-1]\n            outs.append(ndarr)\n \n    if kwargs.get(\"generate_video\", False): # Generate and save video of the entire reverse process. \n        frames2vid(outs, kwargs['save_path'])\n        display(Image.fromarray(outs[-1][:, :, ::-1])) # Display the image at the final timestep of the reverse process.\n        return None\n \n    else: # Display and save the image at the final timestep of the reverse process. \n        x = inverse_transform(x).type(torch.uint8)\n        grid = make_grid(x, nrow=nrow, pad_value=255.0).to(\"cpu\")\n        pil_image = TF.functional.to_pil_image(grid)\n        pil_image.save(kwargs['save_path'], format=save_path[-3:].upper())\n        display(pil_image)\n        for img in x:\n            display(TF.functional.to_pil_image(img))\n        return None","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:16.924596Z","iopub.execute_input":"2023-05-27T18:18:16.925123Z","iopub.status.idle":"2023-05-27T18:18:16.942448Z","shell.execute_reply.started":"2023-05-27T18:18:16.925086Z","shell.execute_reply":"2023-05-27T18:18:16.941468Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import math\nfrom torch.cuda import amp\n\n@dataclass\nclass ModelConfig:\n    BASE_CH = 64  # 64, 128, 256, 256\n    BASE_CH_MULT = (1, 2, 4, 4) # 32, 16, 8, 8 \n    APPLY_ATTENTION = (False, True, True, False)\n    DROPOUT_RATE = 0.1\n    TIME_EMB_MULT = 4 # 128\n \nmodel = UNet(\n    input_channels          = TrainingConfig.IMG_SHAPE[0],\n    output_channels         = TrainingConfig.IMG_SHAPE[0],\n    base_channels           = ModelConfig.BASE_CH,\n    base_channels_multiples = ModelConfig.BASE_CH_MULT,\n    apply_attention         = ModelConfig.APPLY_ATTENTION,\n    dropout_rate            = ModelConfig.DROPOUT_RATE,\n    time_multiple           = ModelConfig.TIME_EMB_MULT,\n)\nmodel.to(BaseConfig.DEVICE)\n \noptimizer = torch.optim.AdamW(model.parameters(), lr=TrainingConfig.LR) # Original → Adam\n \ndataloader = get_dataloader(\n    dataset_name  = BaseConfig.DATASET,\n    batch_size    = TrainingConfig.BATCH_SIZE,\n    device        = BaseConfig.DEVICE,\n    pin_memory    = True,\n    num_workers   = TrainingConfig.NUM_WORKERS,\n)\n \nloss_fn = nn.MSELoss()\n \nsd = SimpleDiffusion(\n    num_diffusion_timesteps = TrainingConfig.TIMESTEPS,\n    img_shape               = TrainingConfig.IMG_SHAPE,\n    device                  = BaseConfig.DEVICE,\n)\n \nscaler = amp.GradScaler() # For mixed-precision training.","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:16.946009Z","iopub.execute_input":"2023-05-27T18:18:16.946416Z","iopub.status.idle":"2023-05-27T18:18:17.291500Z","shell.execute_reply.started":"2023-05-27T18:18:16.946378Z","shell.execute_reply":"2023-05-27T18:18:17.290205Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def setup_log_directory(config):\n    '''Log and Model checkpoint directory Setup'''\n    \n    if os.path.isdir(config.root_log_dir):\n        # Get all folders numbers in the root_log_dir\n        folder_numbers = [int(folder.replace(\"version_\", \"\")) for folder in os.listdir(config.root_log_dir)]\n        \n        # Find the latest version number present in the log_dir\n        last_version_number = max(folder_numbers)\n\n        # New version name\n        version_name = f\"version_{last_version_number + 1}\"\n\n    else:\n        version_name = config.log_dir\n\n    # Update the training config default directory \n    log_dir        = os.path.join(config.root_log_dir,        version_name)\n    checkpoint_dir = os.path.join(config.root_checkpoint_dir, version_name)\n\n    # Create new directory for saving new experiment version\n    os.makedirs(log_dir,        exist_ok=True)\n    os.makedirs(checkpoint_dir, exist_ok=True)\n\n    print(f\"Logging at: {log_dir}\")\n    print(f\"Model Checkpoint at: {checkpoint_dir}\")\n    \n    return log_dir, checkpoint_dir\n\ngenerate_gif = False\ntotal_epochs = TrainingConfig.NUM_EPOCHS + 1\nlog_dir, checkpoint_dir = setup_log_directory(config=BaseConfig())\ngenerate_video = False\next = \".mp4\" if generate_gif else \".png\"","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:17.293558Z","iopub.execute_input":"2023-05-27T18:18:17.294484Z","iopub.status.idle":"2023-05-27T18:18:17.308088Z","shell.execute_reply.started":"2023-05-27T18:18:17.294434Z","shell.execute_reply":"2023-05-27T18:18:17.306702Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# TRAIN","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:17.310234Z","iopub.execute_input":"2023-05-27T18:18:17.311327Z","iopub.status.idle":"2023-05-27T18:18:17.320739Z","shell.execute_reply.started":"2023-05-27T18:18:17.311280Z","shell.execute_reply":"2023-05-27T18:18:17.319686Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\nfrom torchmetrics import MeanMetric\nfrom tqdm import tqdm\n\n# for epoch in range(1, total_epochs):\n#     torch.cuda.empty_cache()\n#     gc.collect()\n     \n#     # Algorithm 1: Training\n#     train_one_epoch(model, sd, dataloader, optimizer, scaler, loss_fn, epoch=epoch)\n \n#     if (epoch == 500):\n#         save_path = os.path.join(log_dir, f\"{epoch}{ext}\")\n        \n#           # clear_output()\n#         checkpoint_dict = {\n#             \"opt\": optimizer.state_dict(),\n#             \"scaler\": scaler.state_dict(),\n#             \"model\": model.state_dict()\n#         }\n#         torch.save(checkpoint_dict, os.path.join(checkpoint_dir, \"ckpt.pt\"))\n#         torch.save(checkpoint_dict, os.path.join(checkpoint_dir, f\"{epoch}-ckpt.pt\"))\n        \n#         # Algorithm 2: Sampling\n#         for iter in range(4):\n#             reverse_diffusion(model, sd, timesteps=TrainingConfig.TIMESTEPS, \n#                               num_images=2, generate_video=generate_video, save_path=save_path, \n#                               img_shape=TrainingConfig.IMG_SHAPE, device=BaseConfig.DEVICE, nrow=4,\n#             )\n#         del checkpoint_dict\n\n# save results","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:17.323989Z","iopub.execute_input":"2023-05-27T18:18:17.324389Z","iopub.status.idle":"2023-05-27T18:18:17.331450Z","shell.execute_reply.started":"2023-05-27T18:18:17.324349Z","shell.execute_reply":"2023-05-27T18:18:17.330143Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nmodel_dir = \"/kaggle/input/diffusionmodels\"\nmodel_files = {\n    \"forest-ckpt.pt\": \"forest\",\n    \"mountain_focused-cpkt.pt\": \"mountain_focused\",\n    \"mountain_smooth-cpkt.pt\": \"mountain_smooth\",\n    \"seaside_dark-ckpt.pt\": \"seaside_dark\",\n    \"seaside_light-ckpt.pt\": \"seaside_light\",\n    \"seaside_unstable-ckpt.pt\": \"seaside_unstable\"\n}\n\nfor model_file, model_name in model_files.items():\n    torch.cuda.empty_cache()\n    checkpoint_path = os.path.join(model_dir, model_file)\n    checkpoint_dict = torch.load(checkpoint_path, map_location=torch.device('cpu'))\n    model.load_state_dict(checkpoint_dict[\"model\"])\n\n    # Algorithm 2: Sampling\n    for iter in range(4):\n        save_path = os.path.join(log_dir, f\"{model_name}_{iter}{ext}\")\n        reverse_diffusion(model, sd, timesteps=TrainingConfig.TIMESTEPS, \n                          num_images=2, generate_video=generate_video, save_path=save_path, \n                          img_shape=TrainingConfig.IMG_SHAPE, device=BaseConfig.DEVICE, nrow=4)","metadata":{"execution":{"iopub.status.busy":"2023-05-27T18:18:17.333066Z","iopub.execute_input":"2023-05-27T18:18:17.334115Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from datetime import datetime\n\ndef to_device(data, device):\n    \"\"\"Move tensor(s) to chosen device\"\"\"\n    if isinstance(data, (list, tuple)):\n        return [to_device(x, device) for x in data]\n    return data.to(device, non_blocking=True)\n\n\ndef frames2vid(images, save_path):\n\n    WIDTH = images[0].shape[1]\n    HEIGHT = images[0].shape[0]\n\n#     fourcc = cv2.VideoWriter_fourcc(*'XVID')\n#     fourcc = 0\n    fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n    video = cv2.VideoWriter(save_path, fourcc, 25, (WIDTH, HEIGHT))\n\n    # Appending the images to the video one by one\n    for image in images:\n        video.write(image)\n\n    # Deallocating memories taken for window creation\n    # cv2.destroyAllWindows()\n    video.release()\n    return \n\ndef display_gif(gif_path):\n    b64 = base64.b64encode(open(gif_path,'rb').read()).decode('ascii')\n    display(HTML(f'<img src=\"data:image/gif;base64,{b64}\" />'))\n\n\nmodel = UNet(\n    input_channels          = TrainingConfig.IMG_SHAPE[0],\n    output_channels         = TrainingConfig.IMG_SHAPE[0],\n    base_channels           = ModelConfig.BASE_CH,\n    base_channels_multiples = ModelConfig.BASE_CH_MULT,\n    apply_attention         = ModelConfig.APPLY_ATTENTION,\n    dropout_rate            = ModelConfig.DROPOUT_RATE,\n    time_multiple           = ModelConfig.TIME_EMB_MULT,\n)\nmodel.load_state_dict(torch.load(os.path.join(checkpoint_dir, \"ckpt.pt\"), map_location='cpu')['model'])\n \nmodel.to(BaseConfig.DEVICE)\n \nsd = SimpleDiffusion(\n    num_diffusion_timesteps = TrainingConfig.TIMESTEPS,\n    img_shape               = TrainingConfig.IMG_SHAPE,\n    device                  = BaseConfig.DEVICE,\n)\n \ngenerate_video = True\n\next = \".mp4\" if generate_video else \".png\"\nfilename = f\"{datetime.now().strftime('%Y%m%d-%H%M%S')}{ext}\"\n\nsave_path = os.path.join(log_dir, filename)\n\n\nreverse_diffusion(\n    model,\n    sd,\n    num_images=2,\n    generate_video=generate_video,\n    save_path=save_path,\n    timesteps=1000,\n    img_shape=TrainingConfig.IMG_SHAPE,\n    device=BaseConfig.DEVICE,\n    nrow=8,\n)\nprint(save_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nos.chdir(r'/kaggle/working')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from IPython.display import FileLink\nFileLink(r'Logs_Checkpoints/Inference/version_7/20230409-153739.png')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ls Logs_Checkpoints/Inference/\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}