First test of backend

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import spaces
 from src.utils import *
+from src.backend import *
 
 theme = gr.themes.Soft(primary_hue="cyan", secondary_hue="zinc", font=[gr.themes.GoogleFont("Source Sans 3", weights=(400, 600)),'arial'])
 
@@ -9,7 +10,9 @@ with gr.Blocks(theme=theme) as demo:
         gr.Markdown("# 🗾 COP-GEN-Beta: Unified Generative Modelling of COPernicus Imagery Thumbnails")
         gr.Markdown("### Miguel Espinosa, Valerio Marsocci, Yuru Jia, Elliot J. Crowley, Mikolaj Czerkawski")
         gr.Markdown('[[Website](https://miquel-espinosa.github.io/cop-gen-beta/)] [[GitHub](https://github.com/miquel-espinosa/COP-GEN-Beta)] [[Model](https://huggingface.co/mespinosami/COP-GEN-Beta)] [[Dataset](https://huggingface.co/Major-TOM)]')
-
+        
+         gr.Markdown('⚠️ NOTE: This is a protoype Beta model of COP-GEN. It is based on image thumbnails of Major TOM and does not yet support raw source data. The hillshade visualisation is used for elevation. The full model COP-GEN is coming soon.')
+    
     with gr.Column(elem_classes="abstract"):
 
         with gr.Accordion("Abstract", open=False) as abstract:
@@ -48,8 +51,7 @@ with gr.Blocks(theme=theme) as demo:
                 dem_output = gr.Image(label="DEM (Elevation)", interactive=False)
 
         with gr.Accordion("Advanced Options", open=False) as advanced_options:
-            num_inference_steps_slider = gr.Slider(minimum=10, maximum=1000, step=10, value=50, label="Inference Steps")
-            guidance_scale_slider = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, value=7.5, label="Guidance Scale")
+            num_inference_steps_slider = gr.Slider(minimum=10, maximum=1000, step=10, value=10, label="Inference Steps")
             with gr.Row():
                 seed_number = gr.Number(value=6378, label="Seed")
                 seed_checkbox = gr.Checkbox(value=True, label="Random")
@@ -61,8 +63,12 @@ with gr.Blocks(theme=theme) as demo:
             
         generate_button.click(
             fn=generate_output,
-            inputs=[s2l1c_input, s2l2a_input, s1rtc_input, dem_input, num_inference_steps_slider, guidance_scale_slider, seed_number, seed_checkbox],
+            inputs=[s2l1c_input, s2l1c_active,
+                    s2l2a_input, s2l2a_active,
+                    s1rtc_input, s1rtc_active,
+                    dem_input, dem_active,
+                    num_inference_steps_slider, seed_number, seed_checkbox],
             outputs=[s2l1c_output, s2l2a_output, s1rtc_output, dem_output],
         )
 
-demo.queue().launch()
+demo.queue().launch(share=True)

requirements.txt

@@ -8,3 +8,4 @@ scikit-learn
 huggingface_hub
 transformers==4.51.1
 accelerate==1.5.2
+ml_collections

src/backend.py

@@ -0,0 +1,278 @@
+import os
+import torch
+import numpy as np
+from PIL import Image
+import ml_collections
+from torchvision.utils import save_image, make_grid
+import torch.nn.functional as F
+import einops
+import random
+import torchvision.transforms as standard_transforms
+
+from huggingface_hub import hf_hub_download
+hf_hub_download(repo_id="thu-ml/unidiffuser-v1", filename="autoencoder_kl.pth", local_dir='./models')
+hf_hub_download(repo_id="mespinosami/COP-GEN-Beta", filename="nnet_ema_114000.pth", local_dir='./models')
+
+import sys
+sys.path.append('./src/COP-GEN-Beta')
+
+import libs
+from dpm_solver_pp import DPM_Solver, NoiseScheduleVP
+from sample_n_triffuser import set_seed, stable_diffusion_beta_schedule, unpreprocess
+import utils
+
+from diffusers import AutoencoderKL
+from .Triffuser import *
+
+# Function to load model
+def load_model(device='cuda'):
+    nnet = Triffuser(num_modalities=4)
+    checkpoint = torch.load('models/nnet_ema_114000.pth', map_location='cuda')
+    nnet.load_state_dict(checkpoint)
+    nnet.to(device)
+    nnet.eval()
+    
+    autoencoder = libs.autoencoder.get_model(pretrained_path = "models/autoencoder_kl.pth")
+    autoencoder.to(device)
+    autoencoder.eval()
+
+    return nnet, autoencoder
+
+print('Loading COP-GEN-Beta model...')
+nnet, autoencoder = load_model()
+to_PIL = standard_transforms.ToPILImage()
+print('[DONE]')
+
+def get_config(generate_modalities, condition_modalities, seed, num_inference_steps=50):
+    config = ml_collections.ConfigDict()
+    config.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    config.seed = seed
+    config.n_samples = 1
+    config.z_shape = (4, 32, 32)  # Shape of the latent vectors
+    config.sample = {
+        'sample_steps': num_inference_steps,
+        'algorithm': "dpm_solver",
+    }
+    # Model config
+    config.num_modalities = 4  # 4 modalities: DEM, S1RTC, S2L1C, S2L2A
+    config.modalities = ['dem', 's1_rtc', 's2_l1c', 's2_l2a']
+    # Network config
+    config.nnet = {
+        'name': 'triffuser_multi_post_ln',
+        'img_size': 32,
+        'in_chans': 4,
+        'patch_size': 2,
+        'embed_dim': 1024,
+        'depth': 20,
+        'num_heads': 16,
+        'mlp_ratio': 4,
+        'qkv_bias': False,
+        'pos_drop_rate': 0.,
+        'drop_rate': 0.,
+        'attn_drop_rate': 0.,
+        'mlp_time_embed': False,
+        'num_modalities': 4,
+        'use_checkpoint': True,
+    }
+
+    # Parse generate and condition modalities
+    config.generate_modalities = generate_modalities
+    config.generate_modalities = sorted(config.generate_modalities, key=lambda x: config.modalities.index(x))
+    config.condition_modalities = condition_modalities if condition_modalities else []
+    config.condition_modalities = sorted(config.condition_modalities, key=lambda x: config.modalities.index(x))
+    config.generate_modalities_mask = [mod in config.generate_modalities for mod in config.modalities]
+    config.condition_modalities_mask = [mod in config.condition_modalities for mod in config.modalities]
+    # Validate modalities
+    valid_modalities = {'s2_l1c', 's2_l2a', 's1_rtc', 'dem'}
+    for mod in config.generate_modalities + config.condition_modalities:
+        if mod not in valid_modalities:
+            raise ValueError(f"Invalid modality: {mod}. Must be one of {valid_modalities}")
+    # Check that generate and condition modalities don't overlap
+    if set(config.generate_modalities) & set(config.condition_modalities):
+        raise ValueError("Generate and condition modalities must be different")
+    # Default data paths
+    config.nnet_path = 'models/nnet_ema_114000.pth'
+    #config.autoencoder = {"pretrained_path": "assets/stable-diffusion/autoencoder_kl_ema.pth"}
+
+    return config
+
+# Function to prepare image for inference
+def prepare_images(images):
+    transforms = standard_transforms.Compose([
+                   standard_transforms.ToTensor(),
+                   standard_transforms.Normalize(mean=(0.5,), std=(0.5,))
+                 ])
+    img_tensors = []
+    for img in images:
+        img_tensors.append(transforms(img))  # Add batch dimension
+    return img_tensors
+
+
+def run_inference(config, nnet, autoencoder, img_tensors):
+    set_seed(config.seed)
+    img_tensors = [tensor.to(config.device) for tensor in img_tensors]
+    # Create a context tensor for all modalities
+    img_contexts = torch.randn(config.num_modalities, 1, 2 * config.z_shape[0],
+                                config.z_shape[1], config.z_shape[2], device=config.device)
+    with torch.no_grad():
+        # Encode the input images with autoencoder
+        z_conds = [autoencoder.encode_moments(tensor.unsqueeze(0)) for tensor in img_tensors]
+        # Create mapping of conditional modalities indices to the encoded inputs
+        cond_indices = [i for i, is_cond in enumerate(config.condition_modalities_mask) if is_cond]
+        # Check if we have the right number of inputs
+        if len(cond_indices) != len(z_conds):
+            raise ValueError(f"Number of conditioning modalities ({len(cond_indices)}) must match number of input images ({len(z_conds)})")
+        # Assign each encoded input to the corresponding modality
+        for i, z_cond in zip(cond_indices, z_conds):
+            img_contexts[i] = z_cond
+    # Sample values from the distribution (mean and variance)
+    z_imgs = torch.stack([autoencoder.sample(img_context) for img_context in img_contexts])
+    # Generate initial noise for the modalities being generated
+    _z_init = torch.randn(len(config.generate_modalities), 1, *z_imgs[0].shape[1:], device=config.device)
+
+    def combine_joint(z_list):
+        """Combine individual modality tensors into a single concatenated tensor"""
+        return torch.concat([einops.rearrange(z_i, 'B C H W -> B (C H W)') for z_i in z_list], dim=-1)
+
+    def split_joint(x, z_imgs, config):
+        """
+        Split the combined tensor back into individual modality tensors
+        and arrange them according to the full set of modalities
+        """
+        C, H, W = config.z_shape
+        z_dim = C * H * W
+        z_generated = x.split([z_dim] * len(config.generate_modalities), dim=1)
+        z_generated = {modality: einops.rearrange(z_i, 'B (C H W) -> B C H W', C=C, H=H, W=W)
+                    for z_i, modality in zip(z_generated, config.generate_modalities)}
+        z = []
+        for i, modality in enumerate(config.modalities):
+            if modality in config.generate_modalities: # Modalities that are being denoised
+                z.append(z_generated[modality])
+            elif modality in config.condition_modalities: # Modalities that are being conditioned on
+                z.append(z_imgs[i])
+            else: # Modalities that are ignored
+                z.append(torch.randn(x.shape[0], C, H, W, device=config.device))
+
+        return z
+
+    _x_init = combine_joint(_z_init) # Initial tensor for the modalities being generated
+    _betas = stable_diffusion_beta_schedule()
+    N = len(_betas)
+
+    def model_fn(x, t_continuous):
+        t = t_continuous * N
+
+        # Create timesteps for each modality based on the generate mask
+        timesteps = [t if mask else torch.zeros_like(t) for mask in config.generate_modalities_mask]
+        # Split the input into a list of tensors for all modalities
+        z = split_joint(x, z_imgs, config)
+        # Call the network with the right format
+        z_out = nnet(z, t_imgs=timesteps)
+        # Select only the generated modalities for the denoising process
+        z_out_generated = [z_out[i]
+                        for i, modality in enumerate(config.modalities)
+                        if modality in config.generate_modalities]
+        # Combine the outputs back into a single tensor
+        return combine_joint(z_out_generated)
+
+    # Sample using the DPM-Solver with exact parameters from sample_n_triffuser.py
+    noise_schedule = NoiseScheduleVP(schedule='discrete', betas=torch.tensor(_betas, device=config.device).float())
+    dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True, thresholding=False)
+
+    # Generate samples
+    with torch.no_grad():
+        with torch.autocast(device_type=config.device):
+            x = dpm_solver.sample(_x_init, steps=config.sample.sample_steps, eps=1. / N, T=1.)
+
+    # Split the result back into individual modality tensors
+    _zs = split_joint(x, z_imgs, config)
+
+    # Replace conditional modalities with the original images
+    for i, mask in enumerate(config.condition_modalities_mask):
+        if mask:
+            _zs[i] = z_imgs[i]
+
+    # Decode and unprocess the generated samples
+    generated_samples = []
+    for i, modality in enumerate(config.modalities):
+        if modality in config.generate_modalities:
+            sample = autoencoder.decode(_zs[i]) # Decode the latent representation
+            sample = unpreprocess(sample) # Unpreprocess to [0, 1] range
+            generated_samples.append((modality, sample))
+
+    return generated_samples
+
+def custom_inference(images, generate_modalities, condition_modalities, num_inference_steps, seed=None):
+    """
+    Run custom inference with user-specified parameters
+
+    Args:
+        generate_modalities: List of modalities to generate
+        condition_modalities: List of modalities to condition on
+        image_paths: Path to conditioning image or list of paths (ordered to match condition_modalities)
+
+    Returns:
+        Dict mapping modality names to generated tensors
+    """
+    if seed is None:
+        seed = random.randint(0, int(1e8))
+
+    img_tensors = prepare_images(images)
+
+    config = get_config(generate_modalities, condition_modalities, seed=seed)
+    config.sample.sample_steps = num_inference_steps
+    generated_samples = run_inference(config, nnet, autoencoder, img_tensors)
+    results = {modality: tensor for modality, tensor in generated_samples}
+    
+    return results
+
+def generate_output(s2l1c_input, s2l1c_active, s2l2a_input, s2l2a_active, s1rtc_input, s1rtc_active, dem_input, dem_active,num_inference_steps_slider, seed_number, ignore_seed):
+
+    seed = seed_number if not ignore_seed else None
+
+    images=[]
+    condition_modalities=[]
+    if s2l2a_active:
+        images.append(s2l2a_input)
+        condition_modalities.append('s2_l2a')
+    if s2l1c_active:
+        images.append(s2l1c_input)
+        condition_modalities.append('s2_l1c')
+    if s1rtc_active:
+        images.append(s1rtc_input)
+        condition_modalities.append('s1_rtc')
+    if dem_active:
+        images.append(dem_input)
+        condition_modalities.append('dem')
+                                    
+    imgs_out = custom_inference(
+        images=images,
+        generate_modalities=[el for el in ['s2_l2a', 's2_l1c', 's1_rtc', 'dem'] if el not in condition_modalities],
+        condition_modalities=condition_modalities,
+        num_inference_steps=num_inference_steps_slider,
+        seed=seed
+    )
+
+    output = []
+
+    # Collect outputs
+    if s2l1c_active:
+        output.append(s2l1c_input)
+    else:
+        output.append(to_PIL(imgs_out['s2_l1c'][0]))
+    if s2l2a_active:
+        output.append(s2l2a_input)
+    else:
+        output.append(to_PIL(imgs_out['s2_l2a'][0]))
+    if s1rtc_active:
+        output.append(s1rtc_input)
+    else:
+        output.append(to_PIL(imgs_out['s1_rtc'][0]))
+    if dem_active:
+        output.append(dem_input)
+    else:
+        output.append(to_PIL(imgs_out['dem'][0]))
+        
+    return output
+
+    

src/utils.py

@@ -100,8 +100,8 @@ def get_rows(grid_cell):
                l2a_df, l1c_df, rtc_df, and dem_df. It assumes these DataFrames are defined in the scope.
                Each element of the tuple is a Pandas Series representing a row.
     """
-    return l2a_df[l2a_df.grid_cell == grid_cell].iloc[0], \
-           l1c_df[l1c_df.grid_cell == grid_cell].iloc[0], \
+    return l1c_df[l1c_df.grid_cell == grid_cell].iloc[0], \
+           l2a_df[l2a_df.grid_cell == grid_cell].iloc[0], \
            rtc_df[rtc_df.grid_cell == grid_cell].iloc[0], \
            dem_df[dem_df.grid_cell == grid_cell].iloc[0]