Upload 3 files

app.py

@@ -0,0 +1,75 @@
+import gradio as gr
+from pathlib import Path
+import os
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+import requests
+
+# Function to download the model from Google Drive
+def download_file_from_google_drive(id, destination):
+    URL = "https://drive.google.com/uc?export=download"
+    session = requests.Session()
+    response = session.get(URL, params={'id': id}, stream=True)
+    token = get_confirm_token(response)
+
+    if token:
+        params = {'id': id, 'confirm': token}
+        response = session.get(URL, params=params, stream=True)
+
+    save_response_content(response, destination)
+
+def get_confirm_token(response):
+    for key, value in response.cookies.items():
+        if key.startswith('download_warning'):
+            return value
+    return None
+
+def save_response_content(response, destination):
+    CHUNK_SIZE = 32768
+    with open(destination, "wb") as f:
+        for chunk in response.iter_content(CHUNK_SIZE):
+            if chunk:  # filter out keep-alive new chunks
+                f.write(chunk)
+
+# Replace 'YOUR_FILE_ID' with your actual file ID from Google Drive
+file_id = '1WJ33nys02XpPDsMO5uIZFiLqTuAT_iuV'
+destination = 'ema_ckpt_cond.pt'
+download_file_from_google_drive(file_id, destination)
+
+# Preprocessing
+from modules import PaletteModelV2
+from diffusion import Diffusion_cond
+
+device = 'cuda'
+
+model = PaletteModelV2(c_in=2, c_out=1, num_classes=5, image_size=256, true_img_size=64).to(device)
+ckpt = torch.load(destination, map_location=device)
+model.load_state_dict(ckpt)
+
+diffusion = Diffusion_cond(noise_steps=1000, img_size=256, device=device)
+model.eval()
+
+transform_hmi = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Resize((256, 256)),
+    transforms.RandomVerticalFlip(p=1.0),
+    transforms.Normalize(mean=(0.5,), std=(0.5,))
+])
+
+def generate_image(seed_image):
+    seed_image_tensor = transform_hmi(Image.open(seed_image)).reshape(1, 1, 256, 256).to(device)
+    generated_image = diffusion.sample(model, y=seed_image_tensor, labels=None, n=1)
+    generated_image_pil = transforms.ToPILImage()(generated_image.squeeze().cpu())
+    return generated_image_pil
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_image,
+    inputs="file",
+    outputs="image",
+    title="Magnetogram-to-Magnetogram: Generative Forecasting of Solar Evolution",
+    description="Upload a LoS magnetogram and predict how it is going to be in 24 hours."
+)
+
+iface.launch()

diffusion.py

@@ -0,0 +1,139 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Apr 25 14:45:59 2023
+
+@author: pio-r
+"""
+
+import torch
+from tqdm import tqdm
+import torch.nn as nn
+import logging
+import numpy as np
+
+logging.basicConfig(format="%(asctime)s - %(levelname)s: %(message)s", level=logging.INFO, datefmt="%I:%M:%S")
+
+
+class Diffusion_cond:
+    def __init__(self, noise_steps=1000, beta_start=1e-4, beta_end=0.02, img_size=256, img_channel=1, device="cuda"):
+        self.noise_steps = noise_steps # timestesps
+        self.beta_start = beta_start
+        self.beta_end = beta_end
+        self.img_channel = img_channel
+        self.img_size = img_size
+        self.device = device
+
+        self.beta = self.prepare_noise_schedule().to(device)
+        self.alpha = 1. - self.beta
+        self.alphas_prev = torch.cat([torch.tensor([1.0]).to(device), self.alpha[:-1]], dim=0)
+        self.alpha_hat = torch.cumprod(self.alpha, dim=0)
+        self.alphas_cumprod_prev = torch.cat([torch.tensor([1.0]).to(device), self.alpha_hat[:-1]], dim=0)
+        # self.alphas_cumprod_prev = torch.from_numpy(np.append(1, self.alpha_hat[:-1].cpu().numpy())).to(device)
+    def prepare_noise_schedule(self):
+        return torch.linspace(self.beta_start, self.beta_end, self.noise_steps) # linear variance schedule as proposed by Ho et al 2020
+
+    def noise_images(self, x, t):
+        sqrt_alpha_hat = torch.sqrt(self.alpha_hat[t])[:, None, None, None]
+        sqrt_one_minus_alpha_hat = torch.sqrt(1 - self.alpha_hat[t])[:, None, None, None]
+        Ɛ = torch.randn_like(x)
+        return sqrt_alpha_hat * x + sqrt_one_minus_alpha_hat * Ɛ, Ɛ # equation in the paper from Ho et al that describes the noise processs
+
+    def sample_timesteps(self, n):
+        return torch.randint(low=1, high=self.noise_steps, size=(n,))
+
+    def sample(self, model, n, y, labels, cfg_scale=3, eta=1, sampling_mode='ddpm'):
+        logging.info(f"Sampling {n} new images....")
+        model.eval() # evaluation mode
+        with torch.no_grad(): # algorithm 2 from DDPM
+            x = torch.randn((n, self.img_channel, self.img_size, self.img_size)).to(self.device)
+            for i in tqdm(reversed(range(1, self.noise_steps)), position=0): # reverse loop from T to 1
+                t = (torch.ones(n) * i).long().to(self.device) # create timesteps tensor of length n
+                predicted_noise = model(x, y, labels, t)
+                if cfg_scale > 0:
+                    uncond_predicted_noise = model(x, y, None, t)
+                    predicted_noise = torch.lerp(uncond_predicted_noise, predicted_noise, cfg_scale)
+                 
+                
+                alpha = self.alpha[t][:, None, None, None]
+                alpha_hat = self.alpha_hat[t][:, None, None, None] # this is noise, created in one
+                alpha_prev = self.alphas_cumprod_prev[t][:, None, None, None]
+                beta = self.beta[t][:, None, None, None]
+                # SAMPLING adjusted from Stable diffusion
+                sigma = (
+                            eta
+                            * torch.sqrt((1 - alpha_prev) / (1 - alpha_hat)
+                            * (1 - alpha_hat / alpha_prev))
+                        )
+                if i > 1:
+                    noise = torch.randn_like(x)
+                else:
+                    noise = torch.zeros_like(x)
+                # pred_x0 = 1 / torch.sqrt(alpha) * (x - ((1 - alpha) / (torch.sqrt(1 - alpha_hat))) * predicted_noise)
+                pred_x0 = (x - torch.sqrt(1 - alpha_hat) * predicted_noise) / torch.sqrt(alpha_hat)
+                if sampling_mode == 'ddpm':
+                    x = 1 / torch.sqrt(alpha) * (x - ((1 - alpha) / (torch.sqrt(1 - alpha_hat))) * predicted_noise) + torch.sqrt(beta) * noise
+                elif sampling_mode == 'ddim':
+                    noise = torch.randn_like(x)
+                    nonzero_mask = (
+                                        (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+                                    )
+                    x = ( 
+                         torch.sqrt(alpha_prev) * pred_x0 +
+                         torch.sqrt(1 - alpha_prev - sigma ** 2) * predicted_noise +
+                         nonzero_mask * sigma * noise
+                        )
+                else:
+                    print('The sampler {} is not implemented'.format(sampling_mode))
+                    break
+        model.train() # it goes back to training mode
+        # x = (x.clamp(-1, 1) + 1) / 2 # to be in [-1, 1], the plus 1 and the division by 2 is to bring back values to [0, 1]
+        # x = (x * 255).type(torch.uint8) # to bring in valid pixel range
+        return x
+
+mse = nn.MSELoss()
+
+def psnr(input: torch.Tensor, target: torch.Tensor, max_val: float) -> torch.Tensor:
+    r"""Create a function that calculates the PSNR between 2 images.
+
+    PSNR is Peek Signal to Noise Ratio, which is similar to mean squared error.
+    Given an m x n image, the PSNR is:
+
+    .. math::
+
+        \text{PSNR} = 10 \log_{10} \bigg(\frac{\text{MAX}_I^2}{MSE(I,T)}\bigg)
+
+    where
+
+    .. math::
+
+        \text{MSE}(I,T) = \frac{1}{mn}\sum_{i=0}^{m-1}\sum_{j=0}^{n-1} [I(i,j) - T(i,j)]^2
+
+    and :math:`\text{MAX}_I` is the maximum possible input value
+    (e.g for floating point images :math:`\text{MAX}_I=1`).
+
+    Args:
+        input: the input image with arbitrary shape :math:`(*)`.
+        labels: the labels image with arbitrary shape :math:`(*)`.
+        max_val: The maximum value in the input tensor.
+
+    Return:
+        the computed loss as a scalar.
+
+    Examples:
+        >>> ones = torch.ones(1)
+        >>> psnr(ones, 1.2 * ones, 2.) # 10 * log(4/((1.2-1)**2)) / log(10)
+        tensor(20.0000)
+
+    Reference:
+        https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio#Definition
+    """
+    if not isinstance(input, torch.Tensor):
+        raise TypeError(f"Expected torch.Tensor but got {type(target)}.")
+
+    if not isinstance(target, torch.Tensor):
+        raise TypeError(f"Expected torch.Tensor but got {type(input)}.")
+
+    if input.shape != target.shape:
+        raise TypeError(f"Expected tensors of equal shapes, but got {input.shape} and {target.shape}")
+
+    return 10.0 * torch.log10(max_val**2 / mse(input, target))

modules.py

@@ -0,0 +1,243 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Apr 25 14:28:21 2023
+
+@author: pio-r
+"""
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+
+class EMA:
+    def __init__(self, beta):
+        super().__init__()
+        self.beta = beta
+        self.step = 0
+
+    def update_model_average(self, ma_model, current_model):
+        for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
+            old_weight, up_weight = ma_params.data, current_params.data
+            ma_params.data = self.update_average(old_weight, up_weight)
+
+    def update_average(self, old, new):
+        if old is None:
+            return new
+        return old * self.beta + (1 - self.beta) * new
+
+    def step_ema(self, ema_model, model, step_start_ema=2000):
+        if self.step < step_start_ema:
+            self.reset_parameters(ema_model, model)
+            self.step += 1
+            return
+        self.update_model_average(ema_model, model)
+        self.step += 1
+
+    def reset_parameters(self, ema_model, model):
+        ema_model.load_state_dict(model.state_dict())
+
+class SelfAttention(nn.Module):
+    """
+    Pre Layer norm  -> multi-headed tension -> skip connections -> pass it to
+    the feed forward layer (layer-norm -> 2 multiheadattention)
+    """
+    def __init__(self, channels, size):
+        super(SelfAttention, self).__init__()
+        self.channels = channels
+        self.size = size
+        self.mha = nn.MultiheadAttention(channels, 4, batch_first=True)
+        self.ln = nn.LayerNorm([channels])
+        self.ff_self = nn.Sequential(
+            nn.LayerNorm([channels]),
+            nn.Linear(channels, channels),
+            nn.GELU(),
+            nn.Linear(channels, channels),
+        )
+
+    def forward(self, x):
+        x = x.view(-1, self.channels, self.size * self.size).swapaxes(1, 2)
+        x_ln = self.ln(x)
+        attention_value, _ = self.mha(x_ln, x_ln, x_ln)
+        attention_value = attention_value + x
+        attention_value = self.ff_self(attention_value) + attention_value
+        return attention_value.swapaxes(2, 1).view(-1, self.channels, self.size, self.size)
+
+
+class DoubleConv(nn.Module):
+    """
+    Normal convolution block, with 2d convolution -> Group Norm -> GeLU -> convolution -> Group Norm
+    Possibility to add residual connection providing residual=True
+    """
+    def __init__(self, in_channels, out_channels, mid_channels=None, residual=False):
+        super().__init__()
+        self.residual = residual
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            nn.GroupNorm(1, mid_channels),
+            nn.GELU(),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.GroupNorm(1, out_channels),
+        )
+
+    def forward(self, x):
+        if self.residual:
+            return F.gelu(x + self.double_conv(x))
+        else:
+            return self.double_conv(x)
+
+
+class Down(nn.Module):
+    """
+    maxpool reduce size by half -> 2*DoubleConv -> Embedding layer
+    
+    """
+    def __init__(self, in_channels, out_channels, emb_dim=256):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, in_channels, residual=True),
+            DoubleConv(in_channels, out_channels),
+        )
+
+        self.emb_layer = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear( # linear projection to bring the time embedding to the proper dimension
+                emb_dim,
+                out_channels
+            ),
+        )
+
+    def forward(self, x, t):
+        x = self.maxpool_conv(x)
+        emb = self.emb_layer(t)[:, :, None, None].repeat(1, 1, x.shape[-2], x.shape[-1]) # projection
+        return x + emb
+
+
+class Up(nn.Module):
+    """
+    We take the skip connection which comes from the encoder
+    """
+    def __init__(self, in_channels, out_channels, emb_dim=256):
+        super().__init__()
+
+        self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
+        self.conv = nn.Sequential(
+            DoubleConv(in_channels, in_channels, residual=True),
+            DoubleConv(in_channels, out_channels, in_channels // 2),
+        )
+
+        self.emb_layer = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(
+                emb_dim,
+                out_channels
+            ),
+        )
+        
+    def forward(self, x, skip_x, t):
+        x = self.up(x)
+        x = torch.cat([skip_x, x], dim=1)
+        x = self.conv(x)
+        emb = self.emb_layer(t)[:, :, None, None].repeat(1, 1, x.shape[-2], x.shape[-1])
+        return x + emb
+
+class PaletteModelV2(nn.Module):
+    def __init__(self, c_in=1, c_out=1, image_size=64, time_dim=256, device='cuda', latent=False, true_img_size=64, num_classes=None):
+        super(PaletteModelV2, self).__init__()
+
+        # Encoder
+        self.true_img_size = true_img_size
+        self.image_size = image_size
+        self.time_dim = time_dim
+        self.device = device
+        self.inc = DoubleConv(c_in, self.image_size) # Wrap-up for 2 Conv Layers
+        self.down1 = Down(self.image_size, self.image_size*2) # input and output channels
+        # self.sa1 = SelfAttention(self.image_size*2,int( self.true_img_size/2)) # 1st is channel dim, 2nd current image resolution
+        self.down2 = Down(self.image_size*2, self.image_size*4)
+        # self.sa2 = SelfAttention(self.image_size*4, int(self.true_img_size/4))
+        self.down3 = Down(self.image_size*4, self.image_size*4)
+        # self.sa3 = SelfAttention(self.image_size*4, int(self.true_img_size/8))
+        
+        # Bootleneck
+        self.bot1 = DoubleConv(self.image_size*4, self.image_size*8)
+        self.bot2 = DoubleConv(self.image_size*8, self.image_size*8)
+        self.bot3 = DoubleConv(self.image_size*8, self.image_size*4)
+        
+        # Decoder: reverse of encoder
+        self.up1 = Up(self.image_size*8, self.image_size*2)
+        # self.sa4 = SelfAttention(self.image_size*2, int(self.true_img_size/4))
+        self.up2 = Up(self.image_size*4, self.image_size)
+        # self.sa5 = SelfAttention(self.image_size, int(self.true_img_size/2))
+        self.up3 = Up(self.image_size*2, self.image_size)
+        # self.sa6 = SelfAttention(self.image_size, self.true_img_size)
+        self.outc = nn.Conv2d(self.image_size, c_out, kernel_size=1) # projecting back to the output channel dimensions
+        
+        if num_classes is not None:
+            self.label_emb = nn.Embedding(num_classes, time_dim)
+
+        if latent == True:
+            self.latent = nn.Sequential(
+                nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1),
+                nn.LeakyReLU(0.2),
+                nn.MaxPool2d(kernel_size=2, stride=2),
+                nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
+                nn.LeakyReLU(0.2),
+                nn.MaxPool2d(kernel_size=2, stride=2),
+                nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
+                nn.LeakyReLU(0.2),
+                nn.MaxPool2d(kernel_size=2, stride=2),
+                nn.Flatten(),
+                nn.Linear(64 * 8 * 8, 256)).to(device)    
+  
+    def pos_encoding(self, t, channels):
+        """
+        Input noised images and the timesteps. The timesteps will only be
+        a tensor with the integer timesteps values in it
+        """
+        inv_freq = 1.0 /  (
+            10000 
+            ** (torch.arange(0, channels, 2, device=self.device).float() / channels)
+        )
+        pos_enc_a = torch.sin(t.repeat(1, channels // 2) * inv_freq)
+        pos_enc_b = torch.cos(t.repeat(1, channels // 2) * inv_freq)
+        pos_enc = torch.cat([pos_enc_a, pos_enc_b], dim=-1)
+        return pos_enc 
+
+    def forward(self, x, y, lab, t):
+        # Pass the source image through the encoder network
+        t = t.unsqueeze(-1).type(torch.float)
+        t = self.pos_encoding(t, self.time_dim) # Encoding timesteps is HERE, we provide the dimension we want to encode
+
+        
+        if lab is not None:
+            t += self.label_emb(lab)
+        
+        # t += self.latent(y)
+        
+        # Concatenate the source image and reference image
+        x = torch.cat([x, y], dim=1)
+        
+        x1 = self.inc(x)
+        x2 = self.down1(x1, t)
+        # x2 = self.sa1(x2)
+        x3 = self.down2(x2, t)
+        # x3 = self.sa2(x3)
+        x4 = self.down3(x3, t)
+        # x4 = self.sa3(x4)
+
+        x4 = self.bot1(x4)
+        x4 = self.bot2(x4)
+        x4 = self.bot3(x4)
+        
+        x = self.up1(x4, x3, t) # We note that upsampling box that in the skip connections from encoder 
+        # x = self.sa4(x)
+        x = self.up2(x, x2, t)
+        # x = self.sa5(x)
+        x = self.up3(x, x1, t)
+        # x = self.sa6(x)
+        output = self.outc(x)
+
+        return output