1

.gitignore

@@ -11,6 +11,7 @@ Pytorch-Correlation-extension/
 result
 src/
 DINOv2FeatureV6_LocalAtten_s2_154000.pth
+example/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

app.py

@@ -13,12 +13,23 @@ import uuid
 import urllib.request
 import warnings
 from os import path
-
-warnings.filterwarnings("ignore", message="The detected CUDA version .* minor version mismatch")
-warnings.filterwarnings("ignore", message="There are no g\\+\\+ version bounds defined for CUDA version.*")
-warnings.filterwarnings("ignore", category=UserWarning, module="torch.utils.cpp_extension")
-os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
-os.environ.setdefault("MAX_JOBS", "1")
+from progressbar import progressbar
+import gc
+
+# # 1) 完全禁止 PyTorch 调用 NVML（ZeroGPU/MIG 下经常拿不到 NVML 句柄）
+# os.environ.setdefault("PYTORCH_NO_NVML", "1")
+# # 2) 用 cudaMallocAsync 后端，降低碎片/避免旧分配器的 NVML 路径
+# os.environ.setdefault(
+#     "PYTORCH_CUDA_ALLOC_CONF",
+#     "backend:cudaMallocAsync,expandable_segments:True,garbage_collection_threshold:0.9,max_split_size_mb:64"
+# )
+# # （可选）定位更准：同步执行
+# os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
+# warnings.filterwarnings("ignore", message="The detected CUDA version .* minor version mismatch")
+# warnings.filterwarnings("ignore", message="There are no g\\+\\+ version bounds defined for CUDA version.*")
+# warnings.filterwarnings("ignore", category=UserWarning, module="torch.utils.cpp_extension")
+# os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
+# os.environ.setdefault("MAX_JOBS", "1")
 
 import gradio as gr
 import spaces  # ZeroGPU decorator
@@ -53,6 +64,15 @@ DESC = """
 - 参考图 -> `./colormnet_run_<UUID>/input_ref/<视频名不含扩展>/ref.png`
 """
 
+# ----------------- TEMP WORKDIR -----------------
+TEMP_ROOT = path.join(os.getcwd(), "_colormnet_tmp")
+
+def reset_temp_root():
+    """每次运行前清空并重建临时工作目录。"""
+    if path.isdir(TEMP_ROOT):
+        shutil.rmtree(TEMP_ROOT, ignore_errors=True)
+    os.makedirs(TEMP_ROOT, exist_ok=True)
+
 torch.set_grad_enabled(False)
 
 # ----------------- DEBUG (kept) -----------------
@@ -146,7 +166,7 @@ def video_to_dataset_root(video_path: str, dataset_root: str):
     if idx == 0:
         raise RuntimeError("Input video has no readable frames.")
 
-    return subdir, stem, w, h, fps, idx
+    return subdir, path.splitext(path.basename(video_path))[0], w, h, fps, idx
 
 # ---------- place ref image into ref_root/<video_stem>/ref.png ----------
 def ref_to_dataset_root(ref_image_path: str, ref_root: str, video_stem: str):
@@ -186,8 +206,8 @@ def run_pipeline_cuda(bw_video_path: str, ref_image_path: str, user_config: dict
 
     DEVICE = torch.device("cuda")
 
-    # Workspace in CWD
-    base_run_dir = path.join(os.getcwd(), f"colormnet_run_{uuid.uuid4().hex}")
+    # Workspace in TEMP_ROOT
+    base_run_dir = path.join(TEMP_ROOT, f"colormnet_run_{uuid.uuid4().hex}")
     input_video_root = path.join(base_run_dir, "input_video")
     input_ref_root   = path.join(base_run_dir, "input_ref")
     output_dir       = path.join(base_run_dir, "result")
@@ -195,11 +215,11 @@ def run_pipeline_cuda(bw_video_path: str, ref_image_path: str, user_config: dict
     for p in (base_run_dir, input_video_root, input_ref_root, output_dir):
         ensure_clean_dir(p)
 
-    # 1) 抽帧
+    # 1) 抽帧（把抽帧输出到临时目录中）
     vid_subdir, vid_stem, w, h, fps, n_frames = video_to_dataset_root(bw_video_path, input_video_root)
     assert n_frames > 0, "Input video has no frames."
 
-    # 2) 参考图
+    # 2) 参考图（存到临时目录）
     _ = ref_to_dataset_root(ref_image_path, input_ref_root, vid_stem)
 
     # 3) 配置（字段与 main.py 一致；值从 UI 合并）
@@ -224,6 +244,7 @@ def run_pipeline_cuda(bw_video_path: str, ref_image_path: str, user_config: dict
         "save_scores": False,
         "flip": False,
         "size": -1,
+        "reverse": False,
     }
     config = {**default_config, **(user_config or {})}
     config["enable_long_term"] = not config["disable_long_term"]
@@ -232,18 +253,7 @@ def run_pipeline_cuda(bw_video_path: str, ref_image_path: str, user_config: dict
     meta_dataset = DAVISTestDataset_221128_TransColorization_batch(
         input_video_root, imset=input_ref_root, size=config["size"]
     )
-    meta_list = meta_dataset.get_datasets()
-
-    target_reader = None
-    for vr in meta_list:
-        if getattr(vr, "vid_name", None) == vid_stem:
-            target_reader = vr
-            break
-    if target_reader is None:
-        if len(meta_list) == 1:
-            target_reader = meta_list[0]
-        else:
-            raise RuntimeError(f"未在数据集中找到目标视频子目录：{vid_stem}；可用={ [getattr(v, 'vid_name', '?') for v in meta_list] }")
+    meta_loader = meta_dataset.get_datasets()
 
     # 输出路径规则（与 main.py 一致）
     is_youtube = str(config["dataset"]).startswith("Y")
@@ -264,111 +274,109 @@ def run_pipeline_cuda(bw_video_path: str, ref_image_path: str, user_config: dict
     total_process_time = 0.0
     total_frames = 0
 
-    # 6) 推理（逐帧；内部逻辑与 main.py 对齐；保留调试打印）
-    vid_reader = target_reader
-    # Gradio/Spaces 环境禁止子进程：num_workers=0（否则会触发 daemonic processes 错误）
-    loader = DataLoader(vid_reader, batch_size=1, shuffle=False, num_workers=0, pin_memory=True)
-    vid_name = vid_reader.vid_name
-    vid_length = len(loader)
-
-    # 长时记忆触发逻辑：按 main.py 原样（无除零保护）
-    config['enable_long_term_count_usage'] = (
-        config['enable_long_term'] and
-        (vid_length
-            / (config['max_mid_term_frames'] - config['min_mid_term_frames'])
-            * config['num_prototypes'])
-        >= config['max_long_term_elements']
-    )
+    for vid_reader in progressbar(meta_loader, max_value=len(meta_dataset), redirect_stdout=True):
+        # 6) 推理（逐帧；内部逻辑与 main.py 对齐；保留调试打印）
+        # Gradio/Spaces 环境禁止子进程：num_workers=0（否则会触发 daemonic processes 错误）
+        loader = DataLoader(vid_reader, batch_size=1, shuffle=False, num_workers=0, pin_memory=True)
+        vid_name = vid_reader.vid_name
+        vid_length = len(loader)
+
+        # 长时记忆触发逻辑：按 main.py 原样（无除零保护）
+        config['enable_long_term_count_usage'] = (
+            config['enable_long_term'] and
+            (vid_length
+                / (config['max_mid_term_frames'] - config['min_mid_term_frames'])
+                * config['num_prototypes'])
+            >= config['max_long_term_elements']
+        )
 
-    mapper = MaskMapper()
-    processor = InferenceCore(network, config=config)
-    first_mask_loaded = False
+        mapper = MaskMapper()
+        processor = InferenceCore(network, config=config)
+        first_mask_loaded = False
+
+        for ti, data in enumerate(loader):
+            try:
+                with torch.cuda.amp.autocast(enabled=not config["benchmark"]):
+                    rgb = data['rgb'].cuda()[0]
+                    msk = data.get('mask')
+                    if not config['FirstFrameIsNotExemplar']:
+                        msk = msk[:, 1:3, :, :] if msk is not None else None
+
+                    info = data['info']
+                    frame = info['frame'][0]
+                    shape = info['shape']
+                    need_resize = info['need_resize'][0]
+
+                    if debug_shapes:
+                        print(f"[Loop] frame={ti} rgb={tuple(rgb.shape)} "
+                            f"msk={None if msk is None else tuple(msk.shape)}", flush=True)
+
+                    # timing 与 main.py 一致
+                    start = torch.cuda.Event(enable_timing=True)
+                    end = torch.cuda.Event(enable_timing=True)
+                    start.record()
+
+                    if not first_mask_loaded:
+                        if msk is not None:
+                            first_mask_loaded = True
+                        else:
+                            continue
+
+                    if config['flip']:
+                        rgb = torch.flip(rgb, dims=[-1])
+                        msk = torch.flip(msk, dims=[-1]) if msk is not None else None
 
-    for ti, data in enumerate(loader):
-        try:
-            with torch.cuda.amp.autocast(enabled=not config["benchmark"]):
-                rgb = data['rgb'].cuda()[0]
-                msk = data.get('mask')
-                if not config['FirstFrameIsNotExemplar']:
-                    msk = msk[:, 1:3, :, :] if msk is not None else None
+                    if msk is not None:
+                        msk = torch.Tensor(msk[0]).cuda()
+                        if need_resize:
+                            msk = vid_reader.resize_mask(msk.unsqueeze(0))[0]
+                        processor.set_all_labels(list(range(1, 3)))
+                        labels = range(1, 3)
+                    else:
+                        labels = None
+
+                    if config['FirstFrameIsNotExemplar']:
+                        prob = processor.step_AnyExemplar(
+                            rgb,
+                            msk[:1, :, :].repeat(3, 1, 1) if msk is not None else None,
+                            msk[1:3, :, :] if msk is not None else None,
+                            labels,
+                            end=(ti == vid_length - 1)
+                        )
+                    else:
+                        prob = processor.step(rgb, msk, labels, end=(ti == vid_length - 1))
+
+                    if need_resize:
+                        prob = F.interpolate(prob.unsqueeze(1), shape, mode='bilinear', align_corners=False)[:, 0]
 
-                print(rgb.shape, msk.shape)
+                    end.record()
+                    torch.cuda.synchronize()
+                    total_process_time += (start.elapsed_time(end) / 1000.0)
+                    total_frames += 1
 
-                info = data['info']
-                frame = info['frame'][0]
-                shape = info['shape']
-                need_resize = info['need_resize'][0]
+                    if config['flip']:
+                        prob = torch.flip(prob, dims=[-1])
 
-                if debug_shapes:
-                    print(f"[Loop] frame={ti} rgb={tuple(rgb.shape)} "
-                          f"msk={None if msk is None else tuple(msk.shape)}", flush=True)
+                    if debug_shapes:
+                        try:
+                            print(f"[Loop] prob={tuple(prob.shape)}", flush=True)
+                        except Exception:
+                            pass
 
-                # timing 与 main.py 一致
-                start = torch.cuda.Event(enable_timing=True)
-                end = torch.cuda.Event(enable_timing=True)
-                start.record()
+                    if config['save_scores']:
+                        prob = (prob.detach().cpu().numpy() * 255).astype(np.uint8)
 
-                if not first_mask_loaded:
-                    if msk is not None:
-                        first_mask_loaded = True
-                    else:
-                        continue
+                    if config['save_all'] or info['save'][0]:
+                        this_out_path = path.join(out_path, vid_name)
+                        os.makedirs(this_out_path, exist_ok=True)
 
-                if config['flip']:
-                    rgb = torch.flip(rgb, dims=[-1])
-                    msk = torch.flip(msk, dims=[-1]) if msk is not None else None
+                        out_mask_final = lab2rgb_transform_PIL(torch.cat([rgb[:1, :, :], prob], dim=0))
+                        out_mask_final = (out_mask_final * 255).astype(np.uint8)
+                        Image.fromarray(out_mask_final).save(os.path.join(this_out_path, frame[:-4] + '.png'))
 
-                if msk is not None:
-                    msk = torch.Tensor(msk[0]).cuda()
-                    if need_resize:
-                        msk = vid_reader.resize_mask(msk.unsqueeze(0))[0]
-                    processor.set_all_labels(list(range(1, 3)))
-                    labels = range(1, 3)
-                else:
-                    labels = None
-
-                if config['FirstFrameIsNotExemplar']:
-                    prob = processor.step_AnyExemplar(
-                        rgb,
-                        msk[:1, :, :].repeat(3, 1, 1) if msk is not None else None,
-                        msk[1:3, :, :] if msk is not None else None,
-                        labels,
-                        end=(ti == vid_length - 1)
-                    )
-                else:
-                    prob = processor.step(rgb, msk, labels, end=(ti == vid_length - 1))
-
-                if need_resize:
-                    prob = F.interpolate(prob.unsqueeze(1), shape, mode='bilinear', align_corners=False)[:, 0]
-
-                end.record()
-                torch.cuda.synchronize()
-                total_process_time += (start.elapsed_time(end) / 1000.0)
-                total_frames += 1
-
-                if config['flip']:
-                    prob = torch.flip(prob, dims=[-1])
-
-                if debug_shapes:
-                    try:
-                        print(f"[Loop] prob={tuple(prob.shape)}", flush=True)
-                    except Exception:
-                        pass
-
-                if config['save_scores']:
-                    prob = (prob.detach().cpu().numpy() * 255).astype(np.uint8)
-
-                if config['save_all'] or info['save'][0]:
-                    this_out_path = path.join(out_path, vid_name)
-                    os.makedirs(this_out_path, exist_ok=True)
-
-                    out_mask_final = lab2rgb_transform_PIL(torch.cat([rgb[:1, :, :], prob], dim=0))
-                    out_mask_final = (out_mask_final * 255).astype(np.uint8)
-                    Image.fromarray(out_mask_final).save(os.path.join(this_out_path, frame[:-4] + '.png'))
-
-        except Exception as _e:
-            # 保留完整 traceback，方便定位
-            raise RuntimeError("FRAME_ERROR:\n" + traceback.format_exc())
+            except Exception as _e:
+                # 保留完整 traceback，方便定位
+                raise RuntimeError("FRAME_ERROR:\n" + traceback.format_exc())
 
     if total_process_time > 0:
         print(f'Total processing time: {total_process_time}')
@@ -388,22 +396,24 @@ def run_pipeline_cuda(bw_video_path: str, ref_image_path: str, user_config: dict
     colored_mp4 = path.join(base_run_dir, "colored_output.mp4")
     encode_frames_to_video(frames_dir, colored_mp4, fps=fps)
 
-    # 8) 输出视频到 CWD
+    # 8) 输出视频到 CWD（只保留最终文件）
     final_mp4 = path.join(os.getcwd(), "result.mp4")
     shutil.move(colored_mp4, final_mp4)
+
+    # 清理本次 run 的中间目录；(注：上传的原视频/参考帧位于 TEMP_ROOT，将在下次运行开头被 reset_temp_root 清掉)
     shutil.rmtree(base_run_dir, ignore_errors=True)
 
     return final_mp4
 
 # ----------------- GRADIO HANDLERS -----------------
-@spaces.GPU(duration=1200)
+@spaces.GPU(duration=600)
 def gradio_infer(
     debug_shapes,  # 调试开关（保留）
     bw_video, ref_image,
     first_not_exemplar, dataset, split, save_all, benchmark,
     disable_long_term, max_mid, min_mid, max_long,
     num_proto, top_k, mem_every, deep_update,
-    save_scores, flip, size
+    save_scores, flip, size, reverse  # 新增
 ):
     if not torch.cuda.is_available():
         return None, "ZeroGPU 未分配到 GPU，请重试（或检查 Space 硬件是否为 ZeroGPU）。"
@@ -413,21 +423,36 @@ def gradio_infer(
     if ref_image is None:
         return None, "请上传参考图像。"
 
-    # Video path
+    # —— 每次运行先重置临时目录 —— #
+    reset_temp_root()
+
+    # Video path -> 拷贝到临时目录
     if isinstance(bw_video, dict) and "name" in bw_video:
-        bw_video_path = bw_video["name"]
+        src_video_path = bw_video["name"]
     elif isinstance(bw_video, str):
-        bw_video_path = bw_video
+        src_video_path = bw_video
     else:
         return None, "无法读取视频输入。"
 
-    # Ref path
+    tmp_video_ext = path.splitext(src_video_path)[1] or ".mp4"
+    tmp_video_path = path.join(TEMP_ROOT, "input_video" + tmp_video_ext)
+    try:
+        shutil.copy2(src_video_path, tmp_video_path)
+    except Exception as e:
+        return None, f"复制视频到临时目录失败：{e}"
+
+    # Ref path -> 保存/拷贝到临时目录
+    tmp_ref_path = path.join(TEMP_ROOT, "ref.png")
     if isinstance(ref_image, Image.Image):
-        tmp_ref_path = path.join(os.getcwd(), f"ref_{uuid.uuid4().hex}.png")
-        ref_image.save(tmp_ref_path)
-        ref_path = tmp_ref_path
+        try:
+            ref_image.save(tmp_ref_path)
+        except Exception as e:
+            return None, f"保存参考图像到临时目录失败：{e}"
     elif isinstance(ref_image, str):
-        ref_path = ref_image
+        try:
+            shutil.copy2(ref_image, tmp_ref_path)
+        except Exception as e:
+            return None, f"复制参考图像到临时目录失败：{e}"
     else:
         return None, "无法读取参考图像输入。"
 
@@ -467,16 +492,22 @@ def gradio_infer(
         "save_scores": bool(save_scores) if save_scores is not None else default_config["save_scores"],
         "flip": bool(flip) if flip is not None else default_config["flip"],
         "size": int(size) if size is not None else default_config["size"],
+        "reverse": bool(reverse) if reverse is not None else False,
     }
 
     try:
         out_mp4 = run_pipeline_cuda(
-            bw_video_path, ref_path, user_config, debug_shapes=bool(debug_shapes)
+            tmp_video_path, tmp_ref_path, user_config, debug_shapes=bool(debug_shapes)
         )
         return out_mp4, "完成 ✅"
     except subprocess.CalledProcessError as e:
+        # 出错也可以顺手清一下临时目录（可选）
+        try: shutil.rmtree(TEMP_ROOT, ignore_errors=True)
+        except: pass
         return None, f"运行时错误：\n{e}"
     except Exception as e:
+        try: shutil.rmtree(TEMP_ROOT, ignore_errors=True)
+        except: pass
         return None, f"{e}"
 
 # ----------------- UI -----------------
@@ -490,9 +521,20 @@ with gr.Blocks() as demo:
         inp_video = gr.Video(label="黑白视频（mp4/webm/avi）", interactive=True)
         inp_ref = gr.Image(label="参考图像（RGB）", type="pil")
 
+        gr.Examples(
+            label="示例输入",
+            examples=[
+                ["./example/4.mp4", "./example/4.png"],
+            ],
+            inputs=[inp_video, inp_ref],
+            # 不缓存，避免把推理结果当静态示例
+            cache_examples=False,
+        )
+
     with gr.Accordion("高级参数设置（与 main.py 对齐）", open=False):
         with gr.Row():
-            first_not_exemplar = gr.Checkbox(label="FirstFrameIsNotExemplar", value=False)
+            first_not_exemplar = gr.Checkbox(label="FirstFrameIsNotExemplar", value=True)
+            reverse = gr.Checkbox(label="reverse", value=False)
             dataset = gr.Textbox(label="dataset", value="D16_batch")
             split = gr.Textbox(label="split", value="val")
             save_all = gr.Checkbox(label="save_all", value=True)
@@ -524,7 +566,7 @@ with gr.Blocks() as demo:
             first_not_exemplar, dataset, split, save_all, benchmark,
             disable_long_term, max_mid, min_mid, max_long,
             num_proto, top_k, mem_every, deep_update,
-            save_scores, flip, size
+            save_scores, flip, size, reverse   # reverse 已接入
         ],
         outputs=[out_video, status]
     )
@@ -535,4 +577,4 @@ if __name__ == "__main__":
     except Exception as e:
         print(f"[WARN] 预下载权重失败（首次推理会再试）: {e}")
 
-    demo.queue(max_size=32).launch(server_name="0.0.0.0", server_port=7860)
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860)

inference/data/test_datasets.py

@@ -5,24 +5,31 @@ import json
 from inference.data.video_reader import VideoReader_221128_TransColorization
 
 class DAVISTestDataset_221128_TransColorization_batch:
-    def __init__(self, data_root, imset='2017/val.txt', size=-1):
+    def __init__(self, data_root, imset='2017/val.txt', size=-1, args=None):
         self.image_dir = data_root
         self.mask_dir = imset
         self.size_dir = data_root
         self.size = size
 
-        self.vid_list =  [clip_name for clip_name in sorted(os.listdir(data_root)) if clip_name != '.DS_Store']
+        self.vid_list =  [clip_name for clip_name in sorted(os.listdir(data_root)) if clip_name != '.DS_Store' and not clip_name.startswith('.')]
+        self.ref_img_list = [clip_name for clip_name in sorted(os.listdir(imset)) if clip_name != '.DS_Store' and not clip_name.startswith('.')]
+
+        self.args = args
 
         # print(lst, len(lst), self.vid_list, self.vid_list_DAVIS2016, path.join(data_root, 'ImageSets', imset));assert 1==0
 
     def get_datasets(self):
         for video in self.vid_list:
+            if video not in self.ref_img_list:
+                continue
+
             # print(self.image_dir, video, path.join(self.image_dir, video));assert 1==0
             yield VideoReader_221128_TransColorization(video, 
                 path.join(self.image_dir, video), 
                 path.join(self.mask_dir, video),
                 size=self.size,
                 size_dir=path.join(self.size_dir, video),
+                args=self.args
             )
 
     def __len__(self):

inference/data/video_reader.py

@@ -14,7 +14,7 @@ class VideoReader_221128_TransColorization(Dataset):
     """
     This class is used to read a video, one frame at a time
     """
-    def __init__(self, vid_name, image_dir, mask_dir, size=-1, to_save=None, use_all_mask=False, size_dir=None):
+    def __init__(self, vid_name, image_dir, mask_dir, size=-1, to_save=None, use_all_mask=False, size_dir=None, args=None):
         """
         image_dir - points to a directory of jpg images
         mask_dir - points to a directory of png masks
@@ -35,9 +35,10 @@ class VideoReader_221128_TransColorization(Dataset):
         else:
             self.size_dir = size_dir
 
-        self.frames = [img for img in sorted(os.listdir(self.image_dir)) if img.endswith('.jpg') or img.endswith('.png')]
-        self.palette = Image.open(path.join(mask_dir, sorted(os.listdir(mask_dir))[0])).getpalette()
-        self.first_gt_path = path.join(self.mask_dir, sorted(os.listdir(self.mask_dir))[0])
+        flag_reverse = args.getattr('reverse', False) if args is not None else False
+        self.frames = [img for img in sorted(os.listdir(self.image_dir), reverse=flag_reverse) if (img.endswith('.jpg') or img.endswith('.png')) and not img.startswith('.')]
+        self.palette = Image.open(path.join(mask_dir, sorted([msk for msk in os.listdir(mask_dir) if not msk.startswith('.')])[0])).getpalette()
+        self.first_gt_path = path.join(self.mask_dir, sorted([msk for msk in os.listdir(self.mask_dir) if not msk.startswith('.')])[0])
         self.suffix = self.first_gt_path.split('.')[-1]
 
         if size < 0:
@@ -87,8 +88,11 @@ class VideoReader_221128_TransColorization(Dataset):
             mask = mask.resize((img.shape[2], img.shape[1]), Image.BILINEAR)
 
             mask = self.im_transform(mask)
-            mask_ab = mask[1:3,:,:]
-            data['mask'] = mask_ab
+
+            # keep L channel of reference image in case First frame is not exemplar
+            # mask_ab = mask[1:3,:,:]
+            # data['mask'] = mask_ab
+            data['mask'] = mask
 
         info['shape'] = shape
         info['need_resize'] = not (self.size < 0)

inference/inference_core.py

@@ -109,3 +109,102 @@ class InferenceCore:
                 self.last_deep_update_ti = self.curr_ti
                 
         return unpad(pred_prob_with_bg, self.pad)
+
+    def step_AnyExemplar(self, image, msk_lll=None, msk_ab=None, valid_labels=None, end=False, flag_FirstframeIsExemplar=False):
+        # image: 3*H*W
+        # mask: num_objects*H*W or None
+        divide_by = 112 # 16
+        self.curr_ti += 1
+        image, self.pad = pad_divide_by(image, divide_by)
+        image = image.unsqueeze(0) # add the batch dimension
+
+        is_mem_frame = ((self.curr_ti-self.last_mem_ti >= self.mem_every) or (msk_ab is not None)) and (not end)
+        need_segment = (self.curr_ti >= 0) and ((valid_labels is None) or (len(self.all_labels) != len(valid_labels))) if not flag_FirstframeIsExemplar else (self.curr_ti > 0) and ((valid_labels is None) or (len(self.all_labels) != len(valid_labels)))
+        is_deep_update = (
+            (self.deep_update_sync and is_mem_frame) or  # synchronized
+            (not self.deep_update_sync and self.curr_ti-self.last_deep_update_ti >= self.deep_update_every) # no-sync
+        ) and (not end)
+        is_normal_update = (not self.deep_update_sync or not is_deep_update) and (not end)
+
+        key, shrinkage, selection, f16, f8, f4 = self.network.encode_key(image, 
+                                                    need_ek=(self.enable_long_term or need_segment), 
+                                                    need_sk=is_mem_frame)
+        multi_scale_features = (f16, f8, f4)
+
+        # save as memory if needed
+        if msk_ab is not None and not flag_FirstframeIsExemplar:
+            need_segment = True
+            is_deep_update = False
+
+            msk_lll, _ = pad_divide_by(msk_lll, divide_by)
+            msk_lll = msk_lll.unsqueeze(0) # add the batch dimension
+            key_mask, shrinkage_mask, selection_mask, f16_mask, f8_mask, f4_mask = self.network.encode_key(msk_lll, 
+                                            need_ek=(self.enable_long_term or need_segment), 
+                                            need_sk=is_mem_frame)
+
+            msk_ab, _ = pad_divide_by(msk_ab, divide_by)
+            pred_prob_with_bg = msk_ab
+
+
+            self.memory.create_hidden_state(2, key)
+        
+
+            value_mask, hidden_mask = self.network.encode_value(msk_lll, f16_mask, self.memory.get_hidden(), 
+                                    pred_prob_with_bg.unsqueeze(0), is_deep_update=False)
+
+            # save key-value to memory
+            self.memory.add_memory(key_mask, shrinkage_mask, value_mask, self.all_labels, 
+                                    selection=selection_mask if self.enable_long_term else None)
+            self.last_mem_ti = self.curr_ti
+
+            self.last_ti_key = key_mask
+            self.last_ti_value = value_mask
+
+            if is_deep_update:
+                self.memory.set_hidden(hidden_mask)
+                self.last_deep_update_ti = self.curr_ti
+
+        # segment the current frame is needed
+        if need_segment:
+            memory_readout = self.memory.match_memory(key, selection).unsqueeze(0)
+
+            # short term memory 
+            batch, num_objects, value_dim, h, w = self.last_ti_value.shape
+            last_ti_value = self.last_ti_value.flatten(start_dim=1, end_dim=2)
+
+            if not (msk_ab is not None and not flag_FirstframeIsExemplar):
+                memory_value_short, _ = self.network.short_term_attn(key, self.last_ti_key, last_ti_value, None, key.shape[-2:]) 
+                memory_value_short = memory_value_short.permute(1, 2, 0).view(batch, num_objects, value_dim, h, w)
+                memory_readout += memory_value_short 
+            hidden, _, pred_prob_with_bg = self.network.segment(multi_scale_features, memory_readout, 
+                                    self.memory.get_hidden(), h_out=is_normal_update, strip_bg=False)
+            # remove batch dim
+            pred_prob_with_bg = pred_prob_with_bg[0]
+            pred_prob_no_bg = pred_prob_with_bg
+            if is_normal_update:
+                self.memory.set_hidden(hidden)
+        else:
+            pred_prob_no_bg = pred_prob_with_bg = None
+
+        # use the input mask if any
+        if msk_ab is not None and flag_FirstframeIsExemplar:
+            msk_ab, _ = pad_divide_by(msk_ab, divide_by)
+            pred_prob_with_bg = msk_ab
+
+        # save as memory if needed
+        if is_mem_frame:
+            value, hidden = self.network.encode_value(image, f16, self.memory.get_hidden(), 
+                                    pred_prob_with_bg.unsqueeze(0), is_deep_update=is_deep_update)
+
+            self.memory.add_memory(key, shrinkage, value, self.all_labels, 
+                                    selection=selection if self.enable_long_term else None)
+            self.last_mem_ti = self.curr_ti
+
+            self.last_ti_key = key
+            self.last_ti_value = value
+
+            if is_deep_update:
+                self.memory.set_hidden(hidden)
+                self.last_deep_update_ti = self.curr_ti
+                
+        return unpad(pred_prob_with_bg, self.pad)