Made Improvement

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

app.py

@@ -4,7 +4,7 @@ import gradio as gr
 
 from stldm import InferenceHub
 from stldm.config import STLDM_HKO
-from utilspp import resize, gradio_gif
+from utilspp import resize, gradio_gif, gradio_visualize
 
 def nowcasting(file, cfg_str, ensemble_no):
     # Model Setup
@@ -30,30 +30,38 @@ def nowcasting(file, cfg_str, ensemble_no):
         raise ValueError("The input should have at least 5 frames for STLDM to predict")
     x = x[0, -5:]
     
-    out = {}
-    for i in range(ensemble_no):
+    y_pred, mu = Forecastor(input_x=x, include_mu=True)
+    out = {'Deterministic': mu, 'Ensemble 1': y_pred}
+    for i in range(1, ensemble_no):
       y_pred = Forecastor(input_x=x, include_mu=False)
-      out[f'Ensemble {i+1}'] = torch.cat((x, y_pred), dim=0)
+      out[f'Ensemble {i+1}'] = y_pred
 
+    past_frames = gradio_visualize(x)
     figure = gradio_gif(out, len(out['Ensemble 1']))
 
-    return figure
+    return past_frames, figure
 
 
 
 with gr.Blocks() as demo:
-    gr.Markdown("# STLDM official demo for nowcasting")
+    gr.Markdown("# STLDM Official Demo for **HKO-7** Nowcasting")
     gr.Markdown("Please upload the radar sequences with **at least 5 frames** in the format of .npy file, and **STLDM** will predict the future 20 frames based on the past 5 frames.")
-    gr.Markdown('Please refer to [paper](https://arxiv.org/abs/2512.21118) and [code](https://github.com/sqfoo/stldm_official) for more details about STLDM.')
+    gr.Markdown('**Paper** - [STLDM: Spatio-Temporal Latent Diffusion Model for Precipitation Nowcasting](https://arxiv.org/abs/2512.21118)')
+    gr.Markdown('**Code** - [https://github.com/sqfoo/stldm_official](https://github.com/sqfoo/stldm_official)')
 
+    gr.Markdown("## Input Frames")
     file_input = gr.File(label="Upload the input radar squences", file_types=[".npy"])
+    
+    gr.Markdown("## Parameters")
     cfg_str = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="Classifier Free Guidance Scale")
     ensemble_no = gr.Slider(1, 10, value=2, step=1, label="How many ensemble predictions?")
 
-    output = gr.Image(label="Nowcasting Results")
+    gr.Markdown("## Predictions")
+    input_frames = gr.Image(label="Past 5 frames")
+    prediction = gr.Image(label="Evolving Predictions")
     btn = gr.Button("Forecast Now!")
-    btn.click(fn=nowcasting, inputs=[file_input, cfg_str, ensemble_no], outputs=output)
+    btn.click(fn=nowcasting, inputs=[file_input, cfg_str, ensemble_no], outputs=[input_frames, prediction])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)
 

stldm/inference.py

@@ -84,9 +84,9 @@ class InferenceHub:
 
         input_x = input_x.to(self.model.device)
         if include_mu:
-            y_pred, mu = self.model(input_x, includ_mu=include_mu)
+            y_pred, mu = self.model(input_x, include_mu=include_mu)
         else:
-            y_pred = self.model(input_x, includ_mu=include_mu)
+            y_pred = self.model(input_x, include_mu=include_mu)
             mu = None
 
         if mu is not None:

stldm/stldm_hf.py

@@ -588,7 +588,7 @@ class GaussianDiffusion(
         return loss.mean()
 
     @torch.no_grad()
-    def forward(self, input_x, include_mu=False, **kwargs):
+    def forward(self, input_x, include_mu, **kwargs):
         pred, mu = self.predict(input_x, compute_loss=False)
         if include_mu:
             return pred, mu

utilspp.py

@@ -1,4 +1,5 @@
 import torch
+import numpy as np
 import matplotlib.pyplot as plt
 import torch.nn.functional as F
 
@@ -34,75 +35,102 @@ def to_cpu_tensor(*args):
 
 from tempfile import NamedTemporaryFile
 
+plt.style.use(['science', 'no-latex'])
+VIL_COLORS = [[0, 0, 0],
+            [0.30196078431372547, 0.30196078431372547, 0.30196078431372547],
+            [0.1568627450980392, 0.7450980392156863, 0.1568627450980392],
+            [0.09803921568627451, 0.5882352941176471, 0.09803921568627451],
+            [0.0392156862745098, 0.4117647058823529, 0.0392156862745098],
+            [0.0392156862745098, 0.29411764705882354, 0.0392156862745098],
+            [0.9607843137254902, 0.9607843137254902, 0.0],
+            [0.9294117647058824, 0.6745098039215687, 0.0],
+            [0.9411764705882353, 0.43137254901960786, 0.0],
+            [0.6274509803921569, 0.0, 0.0],
+            [0.9058823529411765, 0.0, 1.0]]
+
+VIL_LEVELS = [0.0, 16.0, 31.0, 59.0, 74.0, 100.0, 133.0, 160.0, 181.0, 219.0, 255.0]
+
 """ Visualize function with colorbar and a line seprate input and output """
-def gradio_visualize(sequences, horizontal=5, skip=1, ypos=0):
+def gradio_visualize(sequence):
     '''
-    input: sequences, a list/dict of numpy/torch arrays with shape (B, T, C, H, W) 
+    input: sequences, a list/dict of numpy/torch arrays with shape (T, C, H, W) 
     C is assumed to be 1 and squeezed 
     If batch > 1, only the first sequence will be printed 
     '''        
-    plt.style.use(['science', 'no-latex'])
-    VIL_COLORS = [[0, 0, 0],
-                [0.30196078431372547, 0.30196078431372547, 0.30196078431372547],
-                [0.1568627450980392, 0.7450980392156863, 0.1568627450980392],
-                [0.09803921568627451, 0.5882352941176471, 0.09803921568627451],
-                [0.0392156862745098, 0.4117647058823529, 0.0392156862745098],
-                [0.0392156862745098, 0.29411764705882354, 0.0392156862745098],
-                [0.9607843137254902, 0.9607843137254902, 0.0],
-                [0.9294117647058824, 0.6745098039215687, 0.0],
-                [0.9411764705882353, 0.43137254901960786, 0.0],
-                [0.6274509803921569, 0.0, 0.0],
-                [0.9058823529411765, 0.0, 1.0]]
-
-    VIL_LEVELS = [0.0, 16.0, 31.0, 59.0, 74.0, 100.0, 133.0, 160.0, 181.0, 219.0, 255.0]
-
-    # First pass: compute the vertical height and convert to proper format
-    vertical = 0
-    display_texts = []
-    if (type(sequences) is dict):
-        temp = []
-        for k, v in sequences.items():
-            vertical += int(np.ceil(v.shape[1] / horizontal)) 
-            temp.append(v)
-            display_texts.append(k)            
-        sequences = temp
-    else:
-        for i, sequence in enumerate(sequences):
-            vertical += int(np.ceil(sequence.shape[1] / horizontal))
-            display_texts.append(f'Item {i+1}')
-    sequences = to_cpu_tensor(*sequences)
-    # Plot the sequences   
-    j = 0
-    fig, axes = plt.subplots(vertical, horizontal, figsize=(2*horizontal, 2*vertical), tight_layout=True)
+    
+    fig_size = 3
+    fig, axes = plt.subplots(1, len(sequence), figsize=(fig_size*len(sequence), fig_size), tight_layout=True)
     plt.subplots_adjust(hspace=0.0, wspace=0.0) # tight layout
-    plt.setp(axes, xticks=[], yticks=[])    
-    for k, sequence in enumerate(sequences):
-        # only take the first batch, now seq[0] is the temporal dim
-        sequence = sequence.squeeze() # (T, H, W)
-        
-        ## =================
-        # = labels of time =
-        if k == 0:
-            for i in range(len(sequence)):
-                axes[j, i].set_xlabel(f'$t-{(len(sequence)-i)-1}$', fontsize=16)
-                axes[j, i].xaxis.set_label_position('top') 
-        elif k == len(sequences)-1:
-            for i in range(len(sequence)):
-                axes[j, i].set_xlabel(f'$t+{skip*i+1}$', fontsize=16)
-                axes[j, i].xaxis.set_label_position('bottom')            
-        ## =================        
-        axes[j, 0].set_ylabel(display_texts[k], fontsize=16)
-        for i, frame in enumerate(sequence):
-            j_shift = j + i // horizontal 
-            i_shift = i % horizontal
-            im = axes[j_shift, i_shift].imshow(frame*255, cmap=ListedColormap(VIL_COLORS), \
-                                          norm=BoundaryNorm(VIL_LEVELS, ListedColormap(VIL_COLORS).N))
-        j += int(np.ceil(sequence.shape[0] / horizontal))    
+    plt.setp(axes, xticks=[], yticks=[])
+
+    for i, frame in enumerate(sequence):
+        axes[i].set_xticks([])
+        axes[i].set_yticks([])
+        axes[i].set_xlabel(f'$t-{(len(sequence)-i)-1}$', fontsize=12)
+        frame = frame.squeeze()
+        im = axes[i].imshow(frame*255, cmap=ListedColormap(VIL_COLORS), norm=BoundaryNorm(VIL_LEVELS, ListedColormap(VIL_COLORS).N))
+
+
+
+    # # First pass: compute the vertical height and convert to proper format
+    # vertical = 0
+    # display_texts = []
+    # if (type(sequences) is dict):
+    #     temp = []
+    #     for k, v in sequences.items():
+    #         vertical += int(np.ceil(v.shape[1] / horizontal)) 
+    #         temp.append(v)
+    #         display_texts.append(k)            
+    #     sequences = temp
+    # else:
+    #     for i, sequence in enumerate(sequences):
+    #         vertical += int(np.ceil(sequence.shape[1] / horizontal))
+    #         display_texts.append(f'Item {i+1}')
+    # sequences = to_cpu_tensor(*sequences)
+    # # Plot the sequences   
+    # j = 0
+    # fig, axes = plt.subplots(vertical, horizontal, figsize=(2*horizontal, 2*vertical), tight_layout=True)
+    # plt.subplots_adjust(hspace=0.0, wspace=0.0) # tight layout
+    # plt.setp(axes, xticks=[], yticks=[])
     
-    ## = plot splittin line =
-    if ypos == 0:
-        ypos = 1 - 1 / len(sequences) - 0.017
-    fig.lines.append(Line2D((0, 1), (ypos, ypos), transform=fig.transFigure, ls='--', linewidth=2, color='#444'))
+    # if vertical == 1:
+    #     for k, sequence in enumerate(sequences.values()):
+    #         for i in range(len(sequence)):
+    #             axes[i].set_xticks([])
+    #             axes[i].set_yticks([])
+    #             axes[i].set_xlabel(f'$t-{(len(sequence)-i)-1}$', fontsize=12)
+    #             frame = sequence[i].squeeze()
+    #             im = axes[i].imshow(frame*255, cmap=ListedColormap(VIL_COLORS), \
+    #                                 norm=BoundaryNorm(VIL_LEVELS, ListedColormap(VIL_COLORS).N), animated=True)
+    # else:
+        
+    #     for k, sequence in enumerate(sequences):
+    #         # only take the first batch, now seq[0] is the temporal dim
+    #         sequence = sequence.squeeze() # (T, H, W)
+            
+    #         ## =================
+    #         # = labels of time =
+    #         if k == 0:
+    #             for i in range(len(sequence)):
+    #                 axes[j, i].set_xlabel(f'$t-{(len(sequence)-i)-1}$', fontsize=16)
+    #                 axes[j, i].xaxis.set_label_position('top') 
+    #         elif k == len(sequences)-1:
+    #             for i in range(len(sequence)):
+    #                 axes[j, i].set_xlabel(f'$t+{skip*i+1}$', fontsize=16)
+    #                 axes[j, i].xaxis.set_label_position('bottom')            
+    #         ## =================        
+    #         axes[j, 0].set_ylabel(display_texts[k], fontsize=16)
+    #         for i, frame in enumerate(sequence):
+    #             j_shift = j + i // horizontal 
+    #             i_shift = i % horizontal
+    #             im = axes[j_shift, i_shift].imshow(frame*255, cmap=ListedColormap(VIL_COLORS), \
+    #                                         norm=BoundaryNorm(VIL_LEVELS, ListedColormap(VIL_COLORS).N))
+    #         j += int(np.ceil(sequence.shape[0] / horizontal))    
+        
+    # # ## = plot splittin line =
+    # # if ypos == 0:
+    # #     ypos = 1 - 1 / len(sequences) - 0.017
+    # # fig.lines.append(Line2D((0, 1), (ypos, ypos), transform=fig.transFigure, ls='--', linewidth=2, color='#444'))
     # color bar
     cax = fig.add_axes([1, 0.05, 0.02, 0.5])
     fig.colorbar(im, cax=cax)
@@ -124,20 +152,20 @@ def gradio_gif(sequences, T):
     C is assumed to be 1 and squeezed 
     If batch > 1, only the first sequence will be printed 
     '''        
-    plt.style.use(['science', 'no-latex'])
-    VIL_COLORS = [[0, 0, 0],
-                [0.30196078431372547, 0.30196078431372547, 0.30196078431372547],
-                [0.1568627450980392, 0.7450980392156863, 0.1568627450980392],
-                [0.09803921568627451, 0.5882352941176471, 0.09803921568627451],
-                [0.0392156862745098, 0.4117647058823529, 0.0392156862745098],
-                [0.0392156862745098, 0.29411764705882354, 0.0392156862745098],
-                [0.9607843137254902, 0.9607843137254902, 0.0],
-                [0.9294117647058824, 0.6745098039215687, 0.0],
-                [0.9411764705882353, 0.43137254901960786, 0.0],
-                [0.6274509803921569, 0.0, 0.0],
-                [0.9058823529411765, 0.0, 1.0]]
-
-    VIL_LEVELS = [0.0, 16.0, 31.0, 59.0, 74.0, 100.0, 133.0, 160.0, 181.0, 219.0, 255.0]
+    # plt.style.use(['science', 'no-latex'])
+    # VIL_COLORS = [[0, 0, 0],
+    #             [0.30196078431372547, 0.30196078431372547, 0.30196078431372547],
+    #             [0.1568627450980392, 0.7450980392156863, 0.1568627450980392],
+    #             [0.09803921568627451, 0.5882352941176471, 0.09803921568627451],
+    #             [0.0392156862745098, 0.4117647058823529, 0.0392156862745098],
+    #             [0.0392156862745098, 0.29411764705882354, 0.0392156862745098],
+    #             [0.9607843137254902, 0.9607843137254902, 0.0],
+    #             [0.9294117647058824, 0.6745098039215687, 0.0],
+    #             [0.9411764705882353, 0.43137254901960786, 0.0],
+    #             [0.6274509803921569, 0.0, 0.0],
+    #             [0.9058823529411765, 0.0, 1.0]]
+
+    # VIL_LEVELS = [0.0, 16.0, 31.0, 59.0, 74.0, 100.0, 133.0, 160.0, 181.0, 219.0, 255.0]
 
     horizontal = len(sequences)
     fig_size = 3
@@ -146,25 +174,25 @@ def gradio_gif(sequences, T):
     plt.setp(axes, xticks=[], yticks=[])
 
     if horizontal == 1:
-        for i, sequence in enumerate(sequences.values()):
+        for i, (key, sequence) in enumerate(sequences.items()):
             axes.set_xticks([])
             axes.set_yticks([])
-            axes.set_xlabel(f'Ensemble {i+1}', fontsize=12)
+            axes.set_xlabel(f'{key}', fontsize=12)
             frame = sequence[0].squeeze()
             im = axes.imshow(frame*255, cmap=ListedColormap(VIL_COLORS), \
                                     norm=BoundaryNorm(VIL_LEVELS, ListedColormap(VIL_COLORS).N), animated=True)
     else:
-        for i, sequence in enumerate(sequences.values()):
+        for i, (key, sequence) in enumerate(sequences.items()):
             axes[i].set_xticks([])
             axes[i].set_yticks([])
-            axes[i].set_xlabel(f'Ensemble {i+1}', fontsize=12)
+            axes[i].set_xlabel(f'{key}', fontsize=12)
             frame = sequence[0].squeeze()
             im = axes[i].imshow(frame*255, cmap=ListedColormap(VIL_COLORS), \
                                     norm=BoundaryNorm(VIL_LEVELS, ListedColormap(VIL_COLORS).N), animated=True)
 
     title = fig.suptitle('', y=0.9, x=0.505, fontsize=16) # Initialize an empty super title
 
-    fig.colorbar(im)
+    # fig.colorbar(im)
 
     def animate(t):
         if horizontal == 1: