Trying to resolve

app.py

@@ -4,8 +4,7 @@ import gradio as gr
 
 from stldm import InferenceHub
 from stldm.config import STLDM_HKO
-from data.dutils import resize
-from utilspp import gradio_visualize, gradio_gif
+from utilspp import resize, gradio_gif
 
 def nowcasting(file, cfg_str, ensemble_no):
     # Model Setup

requirements.txt

@@ -1,147 +1,64 @@
-absl-py==2.0.0
-antlr4-python3-runtime==4.9.3
-anyio==4.12.0
-argon2-cffi==25.1.0
-argon2-cffi-bindings==25.1.0
-arrow==1.4.0
-asttokens==3.0.1
-async-lru==2.0.5
-attrs==25.4.0
-babel==2.17.0
-beautifulsoup4==4.14.3
-bleach==6.2.0
-cachetools==5.3.2
-certifi==2023.11.17
-cffi==2.0.0
-charset-normalizer==3.3.2
-comm==0.2.3
-contourpy==1.3.0
+aiofiles==24.1.0
+annotated-doc==0.0.4
+annotated-types==0.7.0
+anyio==4.12.1
+brotli==1.2.0
+certifi==2026.1.4
+click==8.3.1
+contourpy==1.3.3
 cycler==0.12.1
-debugpy==1.8.17
-decorator==5.2.1
-defusedxml==0.7.1
-einops==0.8.1
-exceptiongroup==1.3.1
-executing==2.2.1
-fastjsonschema==2.21.2
-fonttools==4.45.0
-fqdn==1.5.1
-google-auth==2.23.4
-google-auth-oauthlib==0.4.6
-grpcio
+einops==0.8.2
+fastapi==0.129.0
+ffmpy==1.0.0
+filelock==3.24.3
+fonttools==4.61.1
+fsspec==2026.2.0
+gradio==6.6.0
+gradio_client==2.1.0
+groovy==0.1.2
 h11==0.16.0
-h5py==3.7.0
+hf-xet==1.2.0
 httpcore==1.0.9
 httpx==0.28.1
-idna==3.4
-imageio==2.33.0
-importlib-metadata==6.8.0
-importlib_resources==6.5.2
-ipykernel==6.31.0
-ipython==8.18.1
-ipywidgets==8.1.8
-isoduration==20.11.0
-jedi==0.19.2
+huggingface_hub==1.4.1
+idna==3.11
 Jinja2==3.1.6
-joblib==1.3.2
-json5==0.12.1
-jsonpointer==3.0.0
-jsonschema==4.25.1
-jsonschema-specifications==2025.9.1
-jupyter==1.1.1
-jupyter-console==6.6.3
-jupyter-events==0.12.0
-jupyter-lsp==2.3.0
-jupyter_client==8.6.3
-jupyter_core==5.8.1
-jupyter_server==2.17.0
-jupyter_server_terminals==0.5.3
-jupyterlab==4.5.0
-jupyterlab_pygments==0.3.0
-jupyterlab_server==2.28.0
-jupyterlab_widgets==3.0.16
-kiwisolver==1.4.5
-lark==1.3.1
-lpips==0.1.4
-Markdown==3.5.1
-MarkupSafe==2.1.3
-matplotlib==3.9.4
-matplotlib-inline==0.2.1
-mistune==3.1.4
-nbclient==0.10.2
-nbconvert==7.16.6
-nbformat==5.10.4
-nest-asyncio==1.6.0
-networkx==3.2.1
-notebook==7.5.0
-notebook_shim==0.2.4
-numpy==1.24.4
-oauthlib==3.2.2
-omegaconf==2.3.0
-opencv-python==4.8.0.74
-overrides==7.7.0
-packaging==23.2
-pandas==1.4.3
-pandocfilters==1.5.1
-parso==0.8.5
-pexpect==4.9.0
-Pillow==10.1.0
-platformdirs==4.4.0
-prometheus_client==0.23.1
-prompt_toolkit==3.0.52
-protobuf==3.19.6
-psutil==7.1.3
-ptyprocess==0.7.0
-pure_eval==0.2.3
-pyasn1==0.5.1
-pyasn1-modules==0.3.0
-pycparser==2.23
+kiwisolver==1.4.9
+markdown-it-py==4.0.0
+MarkupSafe==3.0.3
+matplotlib==3.10.8
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.6.1
+numpy==1.26.4
+orjson==3.11.7
+packaging==26.0
+pandas==3.0.1
+pillow==12.1.1
+pydantic==2.12.5
+pydantic_core==2.41.5
+pydub==0.25.1
 Pygments==2.19.2
-pyparsing==3.1.1
-python-dateutil==2.8.2
-python-json-logger==4.0.0
-pytz==2023.3.post1
-PyWavelets==1.5.0
-PyYAML==6.0
-pyzmq==27.1.0
-referencing==0.36.2
-requests==2.31.0
-requests-oauthlib==1.3.1
-rfc3339-validator==0.1.4
-rfc3986-validator==0.1.1
-rfc3987-syntax==1.1.0
-rpds-py==0.27.1
-rsa==4.9
-SciencePlots==2.2.0
-scikit-image==0.19.3
-scikit-learn==1.1.2
-scipy==1.9.1
-Send2Trash==1.8.3
-six==1.16.0
-soupsieve==2.8
-stack-data==0.6.3
-tensorboard==2.9.0
-tensorboard-data-server==0.6.1
-tensorboard-plugin-wit==1.8.1
-terminado==0.18.1
-threadpoolctl==3.2.0
-tifffile==2023.9.26
-tinycss2==1.4.0
-tomli==2.3.0
-torch==1.12.1+cu116
-torchmetrics==0.11.0
-torchvision==0.13.1+cu116
-tornado==6.5.2
-tqdm==4.66.1
-traitlets==5.14.3
-typing_extensions==4.8.0
-tzdata==2025.2
-uri-template==1.3.0
-urllib3==2.1.0
-wcwidth==0.2.14
-webcolors==24.11.1
-webencodings==0.5.1
-websocket-client==1.9.0
-Werkzeug==3.0.1
-widgetsnbextension==4.0.15
-zipp==3.17.0
+pyparsing==3.3.2
+python-dateutil==2.9.0.post0
+python-multipart==0.0.22
+pytz==2025.2
+PyYAML==6.0.3
+rich==14.3.3
+safehttpx==0.1.7
+safetensors==0.7.0
+SciencePlots @ git+https://github.com/garrettj403/SciencePlots@5521f3b8e6c2b15b174bbea82d6662e5bf2c0d7d
+semantic-version==2.10.0
+setuptools==41.2.0
+shellingham==1.5.4
+six==1.17.0
+starlette==0.52.1
+sympy==1.14.0
+tomlkit==0.13.3
+torch==2.2.0
+tqdm==4.67.3
+typer==0.24.0
+typer-slim==0.24.0
+typing-inspection==0.4.2
+typing_extensions==4.15.0
+uvicorn==0.41.0

stldm/stldm.py

@@ -1,8 +1,17 @@
+# constants
+from collections import namedtuple
+from torch.cuda.amp import autocast
+import torch.nn.functional as F
+from einops import reduce
+from tqdm.auto import tqdm
+
 import torch, random
 from torch import nn
 from einops import rearrange
 
 from stldm.submodules import *
+from stldm.modules import SimVPV2_Model, VAE
+
 
 class Down_Block(nn.Module):
     def __init__(self, in_ch, hid_ch, out_ch, time_dim, is_last, patch_size=None, num_groups=8, heads=4, dim_head=32):
@@ -172,12 +181,6 @@ class LDM(nn.Module):
         
         return out
 
-# constants
-from collections import namedtuple
-from torch.cuda.amp import autocast
-import torch.nn.functional as F
-from einops import reduce
-from tqdm.auto import tqdm
 
 ModelPrediction =  namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
 
@@ -583,7 +586,6 @@ class GaussianDiffusion(nn.Module):
         else:
             return pred
 
-from stldm.modules import SimVPV2_Model, VAE
 def model_setup(model_config, print_info=False, cfg_str=None):
     if print_info:
         print('Setup the model with considering temporal attention be (BHW, T, C) ... ...')

stldm/stldm_hf.py

@@ -1,8 +1,17 @@
+# constants
+from collections import namedtuple
+from torch.cuda.amp import autocast
+import torch.nn.functional as F
+from einops import reduce
+from tqdm.auto import tqdm
+
 import torch, random
 from torch import nn
 from einops import rearrange
 
 from stldm.submodules import *
+from stldm.modules import SimVPV2_Model, VAE
+
 
 class Down_Block(nn.Module):
     def __init__(self, in_ch, hid_ch, out_ch, time_dim, is_last, patch_size=None, num_groups=8, heads=4, dim_head=32):
@@ -172,12 +181,7 @@ class LDM(nn.Module):
         
         return out
 
-# constants
-from collections import namedtuple
-from torch.cuda.amp import autocast
-import torch.nn.functional as F
-from einops import reduce
-from tqdm.auto import tqdm
+
 
 ModelPrediction =  namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
 
@@ -591,7 +595,6 @@ class GaussianDiffusion(
         else:
             return pred
 
-from stldm.modules import SimVPV2_Model, VAE
 def model_setup(model_config, print_info=False, cfg_str=None):
     if print_info:
         print('Setup the model with considering temporal attention be (BHW, T, C) ... ...')

stldm/stldm_spatial.py

@@ -1,8 +1,16 @@
+# constants
+from collections import namedtuple
+from torch.cuda.amp import autocast
+import torch.nn.functional as F
+from einops import reduce
+from tqdm.auto import tqdm
+
 import torch, random
 from torch import nn
 from einops import rearrange
 
 from stldm.submodules import *
+from stldm.modules import SimVPV2_Model, VAE
 
 class Down_Block(nn.Module):
     def __init__(self, in_ch, hid_ch, out_ch, time_dim, is_last, patch_size=None, num_groups=8, heads=4, dim_head=32):
@@ -152,12 +160,7 @@ class LDM(nn.Module):
             out = up_block(out, t, hids1.pop(), hids2.pop())
         return out
 
-# constants
-from collections import namedtuple
-from torch.cuda.amp import autocast
-import torch.nn.functional as F
-from einops import reduce
-from tqdm.auto import tqdm
+
 
 ModelPrediction =  namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
 
@@ -562,7 +565,6 @@ class GaussianDiffusion(nn.Module):
         else:
             return pred
 
-from stldm.modules import SimVPV2_Model, VAE
 def model_setup(model_config, print_info=False, cfg_str=None):
     if print_info:
         print('Setup a Spatial diffusion with replacing a Temporal attention with Spatial attention')

utilspp.py

@@ -1,116 +1,17 @@
-import os
 import torch
-import numpy as np
-import lpips as lp
-import pandas as pd
-import torchmetrics
 import matplotlib.pyplot as plt
-from bisect import bisect_right
-import torchvision.transforms as T
-from torch import nn
+import torch.nn.functional as F
 
 from matplotlib.colors import ListedColormap, BoundaryNorm
 from matplotlib.lines import Line2D
+import matplotlib.animation as animation
+import scienceplots
 
-from data import dutils
-
-# =======================================================================
-# Scheduler Helper Function
-# =======================================================================
-
-class SequentialLR(torch.optim.lr_scheduler._LRScheduler):
-    """Receives the list of schedulers that is expected to be called sequentially during
-    optimization process and milestone points that provides exact intervals to reflect
-    which scheduler is supposed to be called at a given epoch.
-
-    Args:
-        schedulers (list): List of chained schedulers.
-        milestones (list): List of integers that reflects milestone points.
-
-    Example:
-        >>> # Assuming optimizer uses lr = 1. for all groups
-        >>> # lr = 0.1     if epoch == 0
-        >>> # lr = 0.1     if epoch == 1
-        >>> # lr = 0.9     if epoch == 2
-        >>> # lr = 0.81    if epoch == 3
-        >>> # lr = 0.729   if epoch == 4
-        >>> scheduler1 = ConstantLR(self.opt, factor=0.1, total_iters=2)
-        >>> scheduler2 = ExponentialLR(self.opt, gamma=0.9)
-        >>> scheduler = SequentialLR(self.opt, schedulers=[scheduler1, scheduler2], milestones=[2])
-        >>> for epoch in range(100):
-        >>>     train(...)
-        >>>     validate(...)
-        >>>     scheduler.step()
-    """
-
-    def __init__(self, optimizer, schedulers, milestones, last_epoch=-1, verbose=False):
-        for scheduler_idx in range(1, len(schedulers)):
-            if (schedulers[scheduler_idx].optimizer != schedulers[0].optimizer):
-                raise ValueError(
-                    "Sequential Schedulers expects all schedulers to belong to the same optimizer, but "
-                    "got schedulers at index {} and {} to be different".format(0, scheduler_idx)
-                )
-        if (len(milestones) != len(schedulers) - 1):
-            raise ValueError(
-                "Sequential Schedulers expects number of schedulers provided to be one more "
-                "than the number of milestone points, but got number of schedulers {} and the "
-                "number of milestones to be equal to {}".format(len(schedulers), len(milestones))
-            )
-        self.optimizer = optimizer
-        self._schedulers = schedulers
-        self._milestones = milestones
-        self.last_epoch = last_epoch + 1
-
-    def step(self, ref=None):
-        self.last_epoch += 1
-        idx = bisect_right(self._milestones, self.last_epoch)
-        if idx > 0 and self._milestones[idx - 1] == self.last_epoch:
-            self._schedulers[idx].step(0)
-        else:
-            # Check HERE
-            if isinstance(self._schedulers[idx], torch.optim.lr_scheduler.ReduceLROnPlateau):
-                self._schedulers[idx].step(ref)
-            else:
-                self._schedulers[idx].step()
-
-    def state_dict(self):
-        """Returns the state of the scheduler as a :class:`dict`.
-
-        It contains an entry for every variable in self.__dict__ which
-        is not the optimizer.
-        The wrapped scheduler states will also be saved.
-        """
-        state_dict = {key: value for key, value in self.__dict__.items() if key not in ('optimizer', '_schedulers')}
-        state_dict['_schedulers'] = [None] * len(self._schedulers)
-
-        for idx, s in enumerate(self._schedulers):
-            state_dict['_schedulers'][idx] = s.state_dict()
-
-        return state_dict
-
-    def load_state_dict(self, state_dict):
-        """Loads the schedulers state.
-
-        Args:
-            state_dict (dict): scheduler state. Should be an object returned
-                from a call to :meth:`state_dict`.
-        """
-        _schedulers = state_dict.pop('_schedulers')
-        self.__dict__.update(state_dict)
-        # Restore state_dict keys in order to prevent side effects
-        # https://github.com/pytorch/pytorch/issues/32756
-        state_dict['_schedulers'] = _schedulers
-
-        for idx, s in enumerate(_schedulers):
-            self._schedulers[idx].load_state_dict(s)
-
-def warmup_lambda(warmup_steps, min_lr_ratio=0.1):
-    def ret_lambda(epoch):
-        if epoch <= warmup_steps:
-            return min_lr_ratio + (1.0 - min_lr_ratio) * epoch / warmup_steps
-        else:
-            return 1.0
-    return ret_lambda
+def resize(seq, size):
+    # seq shape : (B, T, 1, H, W)
+    seq = F.interpolate(seq.squeeze(dim=2), size=size, mode='bilinear', align_corners=False) # (B, T, H, W)
+    seq = seq.clamp(0,1)
+    return seq.unsqueeze(2) # (B, T, 1, H, W)
 
 # =======================================================================
 # Utils in utils :)
@@ -131,281 +32,6 @@ def to_cpu_tensor(*args):
         return out[0]
     return out
 
-def merge_leading_dims(tensor, n=2):
-    '''
-    Merge the first N dimension of a tensor
-    '''
-    return tensor.reshape((-1, *tensor.shape[n:]))
-
-# =======================================================================
-# Model Preparation, saving & loading (copied from utils.py)
-# =======================================================================
-def build_model_name(model_type, model_config):
-    '''
-    Build the model name (without extension)
-    '''
-    model_name = model_type + '_'
-    for k, v in model_config.items():
-        model_name += k
-        if type(v) is list or type(v) is tuple:
-            model_name += '-'
-            for i, item in enumerate(v):
-                model_name += (str(item) if type(item) is not bool else '') + ('-' if i < len(v)-1 else '')                
-        else:
-            model_name += (('-' + str(v)) if type(v) is not bool else '')
-        model_name += '_'
-    return model_name[:-1]
-
-def build_model_path(base_dir, dataset_type, model_type, timestamp=None):
-    if timestamp is None:
-        return os.path.join(base_dir, dataset_type, model_type)
-    elif timestamp == True:
-        return os.path.join(base_dir, dataset_type, model_type, pd.Timestamp.now().strftime('%Y%m%d%H%M%S'))
-    return os.path.join(base_dir, dataset_type, model_type, timestamp)
-
-# =======================================================================
-# Preprocess Function for Loading HKO-7 dataset
-# =======================================================================
-
-def hko7_preprocess(x_seq, x_mask, dt_clip, args):
-    resize = args.resize if 'resize' in args else x_seq.shape[-1]
-    seq_len = args.seq_len if 'seq_len' in args else 5
-
-    # post-process on HKO-10
-    x_seq = x_seq.transpose((1, 0, 2, 3, 4)) / 255. # => (batch_size, seq_length, 1, 480, 480)
-    if 'scale' in args and args.scale == 'non-linear':
-        x_seq = dutils.linear_to_nonlinear_batched(x_seq, dt_clip)
-    else:
-        x_seq = dutils.nonlinear_to_linear_batched(x_seq, dt_clip)
-
-    b, t, c, h, w = x_seq.shape
-    assert c == 1, f'# channels ({c}) != 1'
-
-    # resize (downsample) the images if necessary
-    x_seq = torch.Tensor(x_seq).float().reshape((b*t, c, h, w))
-    if resize != h:
-        tform = T.Compose([
-            T.ToPILImage(), 
-            T.Resize(resize),
-            T.ToTensor(),
-        ])
-    else:
-        tform = T.Compose([])
-
-    x_seq = torch.stack([tform(x_frame) for x_frame in x_seq], dim=0)
-    x_seq = x_seq.reshape((b, t, c, resize, resize))
-
-    x, y = x_seq[:, :seq_len], x_seq[:, seq_len:]
-    return x, y
-
-# =======================================================================
-# Evaluation Metrics-Related 
-# =======================================================================
-
-mae = lambda *args: torch.nn.functional.l1_loss(*args).cpu().detach().numpy()
-mse = lambda *args: torch.nn.functional.mse_loss(*args).cpu().detach().numpy()
-
-def ssim(y_pred, y):
-    y, y_pred = to_cpu_tensor(y, y_pred)
-    b, t, c, h, w = y.shape
-    y = y.reshape((b*t, c, h, w))
-    y_pred = y_pred.reshape((b*t, c, h, w))
-    # to further ensure any of the input is not negative
-    y = torch.clamp(y, 0, 1)
-    y_pred = torch.clamp(y_pred, 0, 1)
-    return torchmetrics.image.StructuralSimilarityIndexMeasure(data_range=1.0)(y_pred, y)
-
-def psnr(y_pred, y):
-    y, y_pred = to_cpu_tensor(y, y_pred)
-    b, t, c, h, w = y.shape
-    y = y.reshape((b*t, c, h, w))
-    y_pred = y_pred.reshape((b*t, c, h, w))
-    acc_score = 0
-    for i in range(b*t):
-        acc_score += torchmetrics.image.PeakSignalNoiseRatio(data_range=1.0)(y_pred[i], y[i]) / (b*t)
-    return acc_score
-
-GLOBAL_LPIPS_OBJ = None # a static variable
-def lpips64(y_pred, y, net='vgg'):
-    # convert the image range into [-1, 1], assuming the input range to be [0, 1]
-    y = merge_leading_dims(y)
-    y_pred = merge_leading_dims(y_pred)
-
-    y = torch.nn.functional.interpolate(y, (64, 64), mode='bicubic').clamp(0,1)
-    y_pred = torch.nn.functional.interpolate(y_pred, (64, 64), mode='bicubic').clamp(0,1)
-    
-    y = (2 * y - 1)
-    y_pred = (2 * y_pred - 1)
-    global GLOBAL_LPIPS_OBJ
-    if GLOBAL_LPIPS_OBJ is None:
-        GLOBAL_LPIPS_OBJ = lp.LPIPS(net=net).to(y.device)
-    return GLOBAL_LPIPS_OBJ(y_pred, y).mean()
-
-def tfpn(y_pred, y, threshold, radius=1):
-    '''
-    convert to cpu, and merge the first two dimensions
-    '''
-    y = merge_leading_dims(y)
-    y_pred = merge_leading_dims(y_pred)
-    with torch.no_grad():
-        if radius > 1:
-            pool = nn.MaxPool2d(radius)
-            y = pool(y)
-            y_pred = pool(y_pred) 
-        y = torch.where(y >= threshold, 1, 0)
-        y_pred = torch.where(y_pred >= threshold, 1, 0)
-        mat = torchmetrics.functional.confusion_matrix(y_pred, y, task='binary', threshold=threshold)
-        (tn, fp), (fn, tp) = to_cpu_tensor(mat)
-    return tp, tn, fp, fn
-
-def tfpn_pool(y_pred, y, threshold, radius):
-    y_pred = merge_leading_dims(y_pred)
-    y = merge_leading_dims(y)
-    pool = nn.MaxPool2d(radius, stride=radius//4 if radius//4 > 0 else radius) 
-    with torch.no_grad():
-        y = torch.where(y>=threshold, 1, 0).float()
-        y_pred = torch.where(y_pred>=threshold, 1, 0).float()
-        y = pool(y)
-        y_pred = pool(y_pred)
-        mat = torchmetrics.functional.confusion_matrix(y_pred, y, task='binary', threshold=threshold)
-        (tn, fp), (fn, tp) = to_cpu_tensor(mat)
-    return tp, tn, fp, fn
-
-def csi(tp, tn, fp, fn):
-    '''Critical Success Index. The larger the better.'''
-    if (tp + fn + fp) < 1e-7:
-        return 0.
-    return tp / (tp + fn + fp)
-
-def hss(tp, tn, fp, fn):
-    '''Heidke Skill Score. (-inf, 1]. Larger better.'''
-    if (tp+fn)*(fn+tn) + (tp+fp)*(fp+tn) == 0:
-        return 0.
-    return 2 * (tp*tn - fp*fn) / ((tp+fn)*(fn+tn) + (tp+fp)*(fp+tn))
-
-# =======================================================================
-# Data Visualization
-# =======================================================================
-
-def torch_visualize(sequences, savedir=None, horizontal=10, vmin=0, vmax=1):
-    '''
-    input: sequences, a list/dict of numpy/torch arrays with shape (B, T, C, H, W) 
-    C is assumed to be 1 and squeezed 
-    If batch > 1, only the first sequence will be printed 
-    '''        
-    # 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.setp(axes, xticks=[], yticks=[])
-    for k, sequence in enumerate(sequences):
-        # only take the first batch, now seq[0] is the temporal dim
-        sequence = sequence[0].squeeze() # (T, H, W)
-        axes[j, 0].set_ylabel(display_texts[k])
-        for i, frame in enumerate(sequence):
-            j_shift = j + i // horizontal 
-            i_shift = i % horizontal
-            axes[j_shift, i_shift].imshow(frame, vmin=vmin, vmax=vmax, cmap='gray')
-        j += int(np.ceil(sequence.shape[0] / horizontal))    
-    if savedir:
-        plt.savefig(savedir + '' if savedir.endswith('.png') else '.png')
-        plt.close()
-    else:
-        plt.show()
-
-""" Visualize function with colorbar and a line seprate input and output """
-def color_visualize(sequences, savedir='', horizontal=5, skip=1, ypos=0):
-    '''
-    input: sequences, a list/dict of numpy/torch arrays with shape (B, 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)
-    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[0].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)
-    ## =================
-    if savedir:
-        plt.savefig(savedir + '' if len(savedir)>0 else 'out.png')
-        plt.close()
-    else:
-        plt.show()
-
 from tempfile import NamedTemporaryFile
 
 """ Visualize function with colorbar and a line seprate input and output """
@@ -491,7 +117,6 @@ def gradio_visualize(sequences, horizontal=5, skip=1, ypos=0):
     
     return file_path
 
-import matplotlib.animation as animation
 
 def gradio_gif(sequences, T):
     '''
@@ -550,17 +175,4 @@ def gradio_gif(sequences, T):
         file_path = ff.name
     
     plt.close()
-    return file_path
-
-# import matplotlib.pyplot as plt
-# import matplotlib.animation as animation
-# def make_gif(frames, save_path):
-#     fig, ax = plt.subplots(figsize=(4,4))
-#     im = ax.imshow(frames[0].squeeze(), cmap='gray', vmin=0, vmax=1, animated=True)
-#     ax.set_axis_off()
-
-#     def update(i):
-#         im.set_array(frames[i].squeeze())
-#         return im, 
-#     animation_fig = 
-#     animation_fig.save(f"./{save_path}.gif")
+    return file_path