save gpu useness

app.py

@@ -1,10 +1,7 @@
 
 import os
-
-
-
 import gradio as gr
-# gr.set_config(schema_inference=False)
+
 
 import numpy as np
 import torch
@@ -12,46 +9,8 @@ from PIL import Image
 from loguru import logger
 from tqdm import tqdm
 from tools.common_utils import save_video
-from dkt.pipelines.wan_video_new import WanVideoPipeline, ModelConfig
-
-# try:
-#     import gradio_client.utils as _gc_utils
-#     if hasattr(_gc_utils, "get_type"):
-#         _orig_get_type = _gc_utils.get_type
-#         def _get_type_safe(schema):
-#             if not isinstance(schema, dict):
-#                 return "Any"
-#             return _orig_get_type(schema)
-#         _gc_utils.get_type = _get_type_safe
-# except Exception:
-#     pass
-
-# # Additional guard: handle boolean JSON Schemas and parsing errors
-# try:
-#     import gradio_client.utils as _gc_utils
-#     # Wrap the internal _json_schema_to_python_type if present
-#     if hasattr(_gc_utils, "_json_schema_to_python_type"):
-#         _orig_internal = _gc_utils._json_schema_to_python_type
-#         def _json_schema_to_python_type_safe(schema, defs=None):
-#             if isinstance(schema, bool):
-#                 return "Any"
-#             try:
-#                 return _orig_internal(schema, defs)
-#             except Exception:
-#                 return "Any"
-#         _gc_utils._json_schema_to_python_type = _json_schema_to_python_type_safe
-
-#     # Also wrap the public json_schema_to_python_type to be extra defensive
-#     if hasattr(_gc_utils, "json_schema_to_python_type"):
-#         _orig_public = _gc_utils.json_schema_to_python_type
-#         def json_schema_to_python_type_safe(schema):
-#             try:
-#                 return _orig_public(schema)
-#             except Exception:
-#                 return "Any"
-#         _gc_utils.json_schema_to_python_type = json_schema_to_python_type_safe
-# except Exception:
-#     pass
+from dkt.pipelines.pipeline import DKTPipeline, ModelConfig
+
 
 import cv2
 import copy
@@ -59,7 +18,7 @@ import trimesh
 
 from os.path import join
 from tools.depth2pcd import depth2pcd
-from moge.model.v2 import MoGeModel
+# from moge.model.v2 import MoGeModel
 
 
 from tools.eval_utils import transfer_pred_disp2depth, colorize_depth_map
@@ -70,12 +29,18 @@ import tempfile
 import spaces
 
 
-PIPE_1_3B = None
-MOGE_MODULE =  None
+
 #* better for bg:  logs/outs/train/remote/sft-T2SQNet_glassverse_cleargrasp_HISS_DREDS_DREDS_glassverse_interiorverse-4gpus-origin-lora128-1.3B-rgb_depth-w832-h480-Wan2.1-Fun-Control-2025-10-28-23:26:41/epoch-0-20000.safetensors
 PROMPT = 'depth'
 NEGATIVE_PROMPT = ''
 
+height = 480
+width = 832
+window_size = 21
+
+
+DKT_PIPELINE = DKTPipeline()
+
 example_inputs = [
     
     ["examples/1.mp4", "1.3B", 5, 3],
@@ -116,15 +81,6 @@ example_inputs = [
     # ["examples/b68045aa2128ab63d9c7518f8d62eafe.mp4", "1.3B", 5, 3],
 ]
 
-        
-
-
-
-height = 480
-width = 832
-window_size = 21
-
-
 
 def resize_frame(frame, height, width):
     frame = np.array(frame)
@@ -287,173 +243,74 @@ def get_model(model_size):
 
 
 
-@spaces.GPU(duration=120)
-@torch.inference_mode()
+
+
+
 def process_video(
     video_file,
     model_size,
     num_inference_steps,
     overlap
 ):
-    print('process_video called')
-    try:
-        pipe = get_model(model_size)
-        if pipe is None:
-            return None, f"Model {model_size} not initialized. Please restart the application."
-        
-        tmp_video_path = video_file
-        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-
-        # 使用临时目录存储所有文件
-        cur_save_dir = tempfile.mkdtemp(prefix=f'dkt_{timestamp}_{model_size}_')
-        
-        
-        original_filename = f"input_{timestamp}.mp4"
-        dst_path = os.path.join(cur_save_dir, original_filename)
-        shutil.copy2(tmp_video_path, dst_path)
-        origin_frames, input_fps = extract_frames_from_video_file(tmp_video_path)
-        
-        if not origin_frames:
-            return None, "Failed to extract frames from video"
-        
-        logger.info(f"Extracted {len(origin_frames)} frames from video")
+    global height
+    global width
+    global window_size
 
-        
-        original_width, original_height = origin_frames[0].size
-        ROTATE = False 
-        if original_width <  original_height:
-            ROTATE = True
-            origin_frames = [x.transpose(Image.ROTATE_90) for x in origin_frames]
-            tmp = original_width
-            original_width = original_height
-            original_height = tmp
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    cur_save_dir = tempfile.mkdtemp(prefix=f'dkt_{timestamp}_{model_size}_')
+    
 
-        
-                
-        global height
-        global width
-        global window_size
 
-        frames = [resize_frame(frame, height, width) for frame in origin_frames]
-        frame_length = len(frames)
-        if (frame_length - 1) % 4 != 0:
-            new_len = ((frame_length - 1) // 4 + 1) * 4 + 1
-            frames = frames + [copy.deepcopy(frames[-1]) for _ in range(new_len - frame_length)]
+    prediction_result = DKT_PIPELINE(video_file, prompt=PROMPT, \
+        negative_prompt=NEGATIVE_PROMPT,\
+            height=height,width=width,num_inference_steps=num_inference_steps,\
+                overlap=overlap, return_rgb=True)
+    
+    
+    
 
-        
-        control_video = frames
-        video, vae_outs = pipe(
-            prompt=PROMPT,
-            negative_prompt=NEGATIVE_PROMPT,
-            control_video=control_video,
-            height=height,
-            width=width,
-            num_frames=len(control_video),
-            seed=1,
-            tiled=False,
-            num_inference_steps=num_inference_steps,
-            sliding_window_size=window_size,
-            sliding_window_stride=window_size - overlap,
-            cfg_scale=1.0,
-        )
-
-        #* moge process 
-        torch.cuda.empty_cache()
-        processed_video = video[:frame_length]
-
-
-        processed_video = [resize_frame(frame, original_height, original_width) for frame in processed_video]
-        if ROTATE:
-            processed_video = [x.transpose(Image.ROTATE_270) for x in processed_video]
-            origin_frames = [x.transpose(Image.ROTATE_270) for x in origin_frames]
-            
+    frame_length = len(prediction_result['rgb_frames'])
+    vis_pc_num = 4
+    indices = np.linspace(0, frame_length-1, vis_pc_num)
+    indices = np.round(indices).astype(np.int32)
+    
+  
+    pcds = DKT_PIPELINE.prediction2pc_v2(prediction_result['depth_map'], prediction_result['rgb_frames'], indices, return_pcd=True)
+    glb_files = []
 
-        output_filename = f"output_{timestamp}.mp4"
-        output_path = os.path.join(cur_save_dir, output_filename)
-        color_predictions = []
-        if PROMPT == 'depth':
-            prediced_depth_map_np = [np.array(item).astype(np.float32).mean(-1)  for item in processed_video]
-            prediced_depth_map_np = np.stack(prediced_depth_map_np)
-            prediced_depth_map_np = prediced_depth_map_np/ 255.0
-            __min = prediced_depth_map_np.min()
-            __max = prediced_depth_map_np.max()
-            prediced_depth_map_np = (prediced_depth_map_np - __min) / (__max - __min)
-            color_predictions = [colorize_depth_map(item) for item in prediced_depth_map_np]
-        else:
-            color_predictions = processed_video
-        save_video(color_predictions, output_path, fps=input_fps, quality=5)
+    for idx, pcd in enumerate(pcds):
+        points = np.asarray(pcd.points)
+        colors = np.asarray(pcd.colors) if pcd.has_colors() else None
 
 
         
-        # todo, inference MoGe only once
-        frame_num = len(origin_frames)
-        resize_W,resize_H = origin_frames[0].size
-        
-        vis_pc_num = 4
-        indices = np.linspace(0, frame_num-1, vis_pc_num)
-        indices = np.round(indices).astype(np.int32)
-        pc_save_dir = os.path.join(cur_save_dir, 'pointclouds')
-        os.makedirs(pc_save_dir, exist_ok=True)
-
-        glb_files = []
-        moge_device = MOGE_MODULE.device if MOGE_MODULE is not None else torch.device("cuda:0")
-        for idx in tqdm(indices):
-            orgin_rgb_frame = origin_frames[idx]
-            predicted_depth = processed_video[idx]
-
-            # Read the input image and convert to tensor (3, H, W) with RGB values normalized to [0, 1]
-            input_image_np = np.array(orgin_rgb_frame)  # Convert PIL Image to numpy array
-            input_image = torch.tensor(input_image_np / 255, dtype=torch.float32, device=moge_device).permute(2, 0, 1) 
-
-            output = MOGE_MODULE.infer(input_image)
-            #* "dict_keys(['points', 'intrinsics', 'depth', 'mask', 'normal'])"
-            moge_intrinsics = output['intrinsics'].cpu().numpy()
-            moge_mask = output['mask'].cpu().numpy()
-            moge_depth = output['depth'].cpu().numpy()
+        points[:, 2] = -points[:, 2]  
+        points[:, 0] = -points[:, 0]  
         
-            predicted_depth = np.array(predicted_depth)
-            predicted_depth = predicted_depth.mean(-1) / 255.0
-    
-            metric_depth = transfer_pred_disp2depth(predicted_depth, moge_depth, moge_mask)
-            
-            moge_intrinsics[0, 0] *= resize_W 
-            moge_intrinsics[1, 1] *= resize_H
-            moge_intrinsics[0, 2] *= resize_W
-            moge_intrinsics[1, 2] *= resize_H
-            
-            # pcd = depth2pcd(metric_depth, moge_intrinsics, color=cv2.cvtColor(input_image_np, cv2.COLOR_BGR2RGB), input_mask=moge_mask, ret_pcd=True)
-            pcd = depth2pcd(metric_depth, moge_intrinsics, color=input_image_np, input_mask=moge_mask, ret_pcd=True)
-            
-            # pcd.points = o3d.utility.Vector3dVector(np.asarray(pcd.points) * np.array([1, -1, -1], dtype=np.float32))
-            
-            apply_filter = True
-            if apply_filter:
-                cl, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=3.0)
-                pcd = pcd.select_by_index(ind)
 
-            #* save pcd:  o3d.io.write_point_cloud(f'{pc_save_dir}/{timestamp}_{idx:02d}.ply', pcd)
-            points = np.asarray(pcd.points)
-            colors = np.asarray(pcd.colors) if pcd.has_colors() else None
+        glb_filename = os.path.join(cur_save_dir, f'{timestamp}_{idx:02d}.glb')
+        success = create_simple_glb_from_pointcloud(points, colors, glb_filename)
+        if not success:
+            logger.warning(f"Failed to save GLB file: {glb_filename}")
 
+        glb_files.append(glb_filename)
 
-            # ==== 新增：上下翻转点云 ====
-            points[:, 2] = -points[:, 2]  
-            points[:, 0] = -points[:, 0]  
-            # =========================
+    
+    
+    
+    #* save depth predictions video
+    output_filename = f"output_{timestamp}.mp4"
+    output_path = os.path.join(cur_save_dir, output_filename)
 
+    
+    cap = cv2.VideoCapture(video_file)
+    input_fps = cap.get(cv2.CAP_PROP_FPS)
+    cap.release()
 
-            glb_filename = os.path.join(pc_save_dir, f'{timestamp}_{idx:02d}.glb')
-            success = create_simple_glb_from_pointcloud(points, colors, glb_filename)
-            if not success:
-                logger.warning(f"Failed to save GLB file: {glb_filename}")
+    save_video(prediction_result['colored_depth_map'], output_path, fps=input_fps, quality=8)
+    return output_path, glb_files
 
-            glb_files.append(glb_filename)
 
-        return output_path, glb_files
-        
-    except Exception as e:
-        logger.error(f"Error processing video: {str(e)}")
-        return None, f"Error: {str(e)}"
 
 
 
@@ -669,20 +526,16 @@ with gr.Blocks(css=css, title="DKT", head=head_html) as demo:
             )
 
     def on_submit(video_file, model_size, num_inference_steps, overlap):
-        print('on_submit is calling')
         logger.info('on_submit is calling')
-
         if video_file is None:
             return None, None, None, None, None, None, "Please upload a video file"
         
         try:
-
             output_path, glb_files = process_video(
                 video_file, model_size, num_inference_steps, overlap
             )
 
             
-            
             if output_path is None:
                 return None, None, None, None, None, None, glb_files
             
@@ -732,23 +585,12 @@ with gr.Blocks(css=css, title="DKT", head=head_html) as demo:
 if __name__ == '__main__':
     
     #* main code, model and moge model initialization
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    logger.info(f"device = {device}")
-    print(f"device = {device}")
     
-    load_model_1_3b(device=device)
-    load_moge_model(device=device)
-    # torch.cuda.empty_cache()
-    logger.info('model init done!')
-    print('model init done!')
 
-    demo.queue().launch(share = True)
-    
-    # demo.queue(
-    #     api_open=False,
-    # ).launch()
 
+
+    demo.queue().launch(share = True)
     
-    # server_name="0.0.0.0", server_port=7860
+ 
 
 

app_old.py

@@ -0,0 +1,756 @@
+
+import os
+
+
+
+import gradio as gr
+# gr.set_config(schema_inference=False)
+
+import numpy as np
+import torch
+from PIL import Image
+from loguru import logger
+from tqdm import tqdm
+from tools.common_utils import save_video
+from dkt.pipelines.pipeline import WanVideoPipeline, ModelConfig
+
+# try:
+#     import gradio_client.utils as _gc_utils
+#     if hasattr(_gc_utils, "get_type"):
+#         _orig_get_type = _gc_utils.get_type
+#         def _get_type_safe(schema):
+#             if not isinstance(schema, dict):
+#                 return "Any"
+#             return _orig_get_type(schema)
+#         _gc_utils.get_type = _get_type_safe
+# except Exception:
+#     pass
+
+# # Additional guard: handle boolean JSON Schemas and parsing errors
+# try:
+#     import gradio_client.utils as _gc_utils
+#     # Wrap the internal _json_schema_to_python_type if present
+#     if hasattr(_gc_utils, "_json_schema_to_python_type"):
+#         _orig_internal = _gc_utils._json_schema_to_python_type
+#         def _json_schema_to_python_type_safe(schema, defs=None):
+#             if isinstance(schema, bool):
+#                 return "Any"
+#             try:
+#                 return _orig_internal(schema, defs)
+#             except Exception:
+#                 return "Any"
+#         _gc_utils._json_schema_to_python_type = _json_schema_to_python_type_safe
+
+#     # Also wrap the public json_schema_to_python_type to be extra defensive
+#     if hasattr(_gc_utils, "json_schema_to_python_type"):
+#         _orig_public = _gc_utils.json_schema_to_python_type
+#         def json_schema_to_python_type_safe(schema):
+#             try:
+#                 return _orig_public(schema)
+#             except Exception:
+#                 return "Any"
+#         _gc_utils.json_schema_to_python_type = json_schema_to_python_type_safe
+# except Exception:
+#     pass
+
+import cv2
+import copy
+import trimesh
+
+from os.path import join
+from tools.depth2pcd import depth2pcd
+from moge.model.v2 import MoGeModel
+
+
+from tools.eval_utils import transfer_pred_disp2depth, colorize_depth_map
+import glob
+import datetime
+import shutil
+import tempfile
+import spaces
+
+
+PIPE_1_3B = None
+MOGE_MODULE =  None
+#* better for bg:  logs/outs/train/remote/sft-T2SQNet_glassverse_cleargrasp_HISS_DREDS_DREDS_glassverse_interiorverse-4gpus-origin-lora128-1.3B-rgb_depth-w832-h480-Wan2.1-Fun-Control-2025-10-28-23:26:41/epoch-0-20000.safetensors
+PROMPT = 'depth'
+NEGATIVE_PROMPT = ''
+
+example_inputs = [
+    
+    ["examples/1.mp4", "1.3B", 5, 3],
+    ["examples/33.mp4", "1.3B", 5, 3],
+    
+    
+    
+    ["examples/7.mp4", "1.3B", 5, 3],
+    ["examples/8.mp4", "1.3B", 5, 3],
+    ["examples/9.mp4", "1.3B", 5, 3],
+    
+    # ["examples/178db6e89ab682bfc612a3290fec58dd.mp4", "1.3B", 5, 3],
+    ["examples/36.mp4", "1.3B", 5, 3],
+    ["examples/39.mp4", "1.3B", 5, 3],
+    
+    # ["examples/b1f1fa44f414d7731cd7d77751093c44.mp4", "1.3B", 5, 3],
+    
+    ["examples/10.mp4", "1.3B", 5, 3],
+    ["examples/30.mp4", "1.3B", 5, 3],
+    ["examples/3.mp4", "1.3B", 5, 3],
+    
+    ["examples/32.mp4", "1.3B", 5, 3],
+    
+    ["examples/35.mp4", "1.3B", 5, 3],
+    
+    ["examples/40.mp4", "1.3B", 5, 3],
+    ["examples/2.mp4", "1.3B", 5, 3],
+
+    # ["examples/31.mp4", "1.3B", 5, 3],
+    # ["examples/DJI_20250912164311_0007_D.mp4", "1.3B", 5, 3],
+    # ["examples/DJI_20250912163642_0003_D.mp4", "1.3B", 5, 3],
+
+    # ["examples/5.mp4", "1.3B", 5, 3],
+
+    # ["examples/1b0daeb776471c7389b36cee53049417.mp4", "1.3B", 5, 3],
+    # ["examples/8a6dfb8cfe80634f4f77ae9aa830d075.mp4", "1.3B", 5, 3],
+    # ["examples/69230f105ad8740e08d743a8ee11c651.mp4", "1.3B", 5, 3],
+    # ["examples/b68045aa2128ab63d9c7518f8d62eafe.mp4", "1.3B", 5, 3],
+]
+
+        
+
+
+
+height = 480
+width = 832
+window_size = 21
+
+
+
+def resize_frame(frame, height, width):
+    frame = np.array(frame)
+    frame = torch.from_numpy(frame).permute(2, 0, 1).unsqueeze(0).float() / 255.0
+    frame = torch.nn.functional.interpolate(frame, (height, width), mode="bicubic", align_corners=False, antialias=True)
+    frame = (frame.squeeze(0).permute(1, 2, 0).clamp(0, 1) * 255).byte().numpy()
+    frame = Image.fromarray(frame)
+    return frame
+
+
+
+def pmap_to_glb(point_map, valid_mask, frame) -> trimesh.Scene:
+    pts_3d = point_map[valid_mask] * np.array([-1, -1, 1])
+    pts_rgb = frame[valid_mask] 
+
+    # Initialize a 3D scene
+    scene_3d = trimesh.Scene()
+
+    # Add point cloud data to the scene
+    point_cloud_data = trimesh.PointCloud(
+        vertices=pts_3d, colors=pts_rgb
+    )
+    
+    scene_3d.add_geometry(point_cloud_data)
+    return scene_3d
+
+
+
+def create_simple_glb_from_pointcloud(points, colors, glb_filename):
+    try:
+        if len(points) == 0:
+            logger.warning(f"No valid points to create GLB for {glb_filename}")
+            return False
+        
+        if colors is not None:
+            # logger.info(f"Adding colors to GLB: shape={colors.shape}, range=[{colors.min():.3f}, {colors.max():.3f}]")
+            pts_rgb = colors
+        else:
+            logger.info("No colors provided, adding default white colors")
+            pts_rgb = np.ones((len(points), 3))
+        
+        valid_mask = np.ones(len(points), dtype=bool)
+        
+        scene_3d = pmap_to_glb(points, valid_mask, pts_rgb)
+        
+        scene_3d.export(glb_filename)
+        # logger.info(f"Saved GLB file using trimesh: {glb_filename}")
+        
+        return True
+        
+    except Exception as e:
+        logger.error(f"Error creating GLB from pointcloud using trimesh: {str(e)}")
+        return False
+
+
+
+
+
+def extract_frames_from_video_file(video_path):
+    try:
+        cap = cv2.VideoCapture(video_path)
+        frames = []
+        
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        if fps <= 0:
+            fps = 15.0
+        
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame_rgb = Image.fromarray(frame_rgb)
+            frames.append(frame_rgb)
+        
+        cap.release()
+        return frames, fps
+    except Exception as e:
+        logger.error(f"Error extracting frames from {video_path}: {str(e)}")
+        return [], 15.0
+
+
+
+def load_moge_model(device="cuda:0"):
+    global MOGE_MODULE
+    if MOGE_MODULE is not None:
+        return MOGE_MODULE
+    logger.info(f"Loading MoGe model on {device}...")
+    MOGE_MODULE = MoGeModel.from_pretrained('Ruicheng/moge-2-vitl-normal').to(device)
+    return MOGE_MODULE
+
+
+def load_model_1_3b(device="cuda:0"):
+    global PIPE_1_3B
+    
+    if PIPE_1_3B is not None:
+        return PIPE_1_3B
+    
+    logger.info(f"Loading 1.3B model on {device}...")
+    
+    pipe = WanVideoPipeline.from_pretrained(
+        torch_dtype=torch.bfloat16,
+        device=device,
+        model_configs=[
+            ModelConfig(
+                model_id="PAI/Wan2.1-Fun-1.3B-Control",
+                origin_file_pattern="diffusion_pytorch_model*.safetensors",
+                offload_device="cpu",
+            ),
+            ModelConfig(
+                model_id="PAI/Wan2.1-Fun-1.3B-Control",
+                origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth",
+                offload_device="cpu",
+            ),
+            ModelConfig(
+                model_id="PAI/Wan2.1-Fun-1.3B-Control",
+                origin_file_pattern="Wan2.1_VAE.pth",
+                offload_device="cpu",
+            ),
+            ModelConfig(
+                model_id="PAI/Wan2.1-Fun-1.3B-Control",
+                origin_file_pattern="models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",
+                offload_device="cpu",
+            ),
+        ],
+        training_strategy="origin",
+    )
+    
+    
+    lora_config = ModelConfig(
+        model_id="Daniellesry/DKT-Depth-1-3B",
+        origin_file_pattern="dkt-1-3B.safetensors",
+        offload_device="cpu",
+    )
+
+    lora_config.download_if_necessary(use_usp=False)
+    
+    pipe.load_lora(pipe.dit, lora_config.path, alpha=1.0)#todo is it work?
+    pipe.enable_vram_management()
+
+    
+    PIPE_1_3B = pipe
+    
+    return pipe
+
+
+
+    
+
+def get_model(model_size):
+    if model_size == "1.3B":
+        assert PIPE_1_3B is not None, "1.3B model not initialized"
+        return PIPE_1_3B
+    else:
+        raise ValueError(f"Unsupported model size: {model_size}")
+
+
+
+
+
+
+
+
+
+@spaces.GPU(duration=120)
+@torch.inference_mode()
+def process_video(
+    video_file,
+    model_size,
+    num_inference_steps,
+    overlap
+):
+    
+    pipe = get_model(model_size)
+
+    if pipe is None:
+        return None, f"Model {model_size} not initialized. Please restart the application."
+    
+    tmp_video_path = video_file
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    
+    cur_save_dir = tempfile.mkdtemp(prefix=f'dkt_{timestamp}_{model_size}_')
+    origin_frames, input_fps = extract_frames_from_video_file(tmp_video_path)
+    
+    if not origin_frames:
+        return None, "Failed to extract frames from video"
+    
+    logger.info(f"Extracted {len(origin_frames)} frames from video")
+
+    original_width, original_height = origin_frames[0].size
+    ROTATE = False 
+    if original_width <  original_height:
+        ROTATE = True
+        origin_frames = [x.transpose(Image.ROTATE_90) for x in origin_frames]
+        tmp = original_width
+        original_width = original_height
+        original_height = tmp
+
+    
+            
+    global height
+    global width
+    global window_size
+
+    frames = [resize_frame(frame, height, width) for frame in origin_frames]
+    frame_length = len(frames)
+    if (frame_length - 1) % 4 != 0:
+        new_len = ((frame_length - 1) // 4 + 1) * 4 + 1
+        frames = frames + [copy.deepcopy(frames[-1]) for _ in range(new_len - frame_length)]
+
+    
+    control_video = frames
+    video, vae_outs = pipe(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        control_video=control_video,
+        height=height,
+        width=width,
+        num_frames=len(control_video),
+        seed=1,
+        tiled=False,
+        num_inference_steps=num_inference_steps,
+        sliding_window_size=window_size,
+        sliding_window_stride=window_size - overlap,
+        cfg_scale=1.0,
+    )
+
+    #* moge process 
+    torch.cuda.empty_cache()
+    processed_video = video[:frame_length]
+
+
+    processed_video = [resize_frame(frame, original_height, original_width) for frame in processed_video]
+    if ROTATE:
+        processed_video = [x.transpose(Image.ROTATE_270) for x in processed_video]
+        origin_frames = [x.transpose(Image.ROTATE_270) for x in origin_frames]
+        
+    
+    color_predictions = []
+    if PROMPT == 'depth':
+        prediced_depth_map_np = [np.array(item).astype(np.float32).mean(-1)  for item in processed_video]
+        prediced_depth_map_np = np.stack(prediced_depth_map_np)
+        prediced_depth_map_np = prediced_depth_map_np/ 255.0
+        __min = prediced_depth_map_np.min()
+        __max = prediced_depth_map_np.max()
+        prediced_depth_map_np = (prediced_depth_map_np - __min) / (__max - __min)
+        color_predictions = [colorize_depth_map(item) for item in prediced_depth_map_np]
+    else:
+        color_predictions = processed_video
+
+    
+    
+    
+
+    #* required parameters  for MoGe 
+
+    # todo, inference MoGe only once
+    
+    resize_W,resize_H = origin_frames[0].size
+    
+    vis_pc_num = 4
+    indices = np.linspace(0, frame_length-1, vis_pc_num)
+    indices = np.round(indices).astype(np.int32)
+    pc_save_dir = os.path.join(cur_save_dir, 'pointclouds')
+    os.makedirs(pc_save_dir, exist_ok=True)
+
+    glb_files = []
+    moge_device = MOGE_MODULE.device if MOGE_MODULE is not None else torch.device("cuda:0")
+
+    for idx in tqdm(indices):
+        orgin_rgb_frame = origin_frames[idx]
+        predicted_depth = processed_video[idx]
+
+        # Read the input image and convert to tensor (3, H, W) with RGB values normalized to [0, 1]
+        input_image_np = np.array(orgin_rgb_frame)  # Convert PIL Image to numpy array
+        input_image = torch.tensor(input_image_np / 255, dtype=torch.float32, device=moge_device).permute(2, 0, 1) 
+
+        output = MOGE_MODULE.infer(input_image)
+        #* "dict_keys(['points', 'intrinsics', 'depth', 'mask', 'normal'])"
+        moge_intrinsics = output['intrinsics'].cpu().numpy()
+        moge_mask = output['mask'].cpu().numpy()
+        moge_depth = output['depth'].cpu().numpy()
+    
+        predicted_depth = np.array(predicted_depth)
+        predicted_depth = predicted_depth.mean(-1) / 255.0
+
+        metric_depth = transfer_pred_disp2depth(predicted_depth, moge_depth, moge_mask)
+        
+        moge_intrinsics[0, 0] *= resize_W 
+        moge_intrinsics[1, 1] *= resize_H
+        moge_intrinsics[0, 2] *= resize_W
+        moge_intrinsics[1, 2] *= resize_H
+        
+        # pcd = depth2pcd(metric_depth, moge_intrinsics, color=cv2.cvtColor(input_image_np, cv2.COLOR_BGR2RGB), input_mask=moge_mask, ret_pcd=True)
+        pcd = depth2pcd(metric_depth, moge_intrinsics, color=input_image_np, input_mask=moge_mask, ret_pcd=True)
+        
+        # pcd.points = o3d.utility.Vector3dVector(np.asarray(pcd.points) * np.array([1, -1, -1], dtype=np.float32))
+        
+        apply_filter = True
+        if apply_filter:
+            cl, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=3.0)
+            pcd = pcd.select_by_index(ind)
+
+        #* save pcd:  o3d.io.write_point_cloud(f'{pc_save_dir}/{timestamp}_{idx:02d}.ply', pcd)
+        points = np.asarray(pcd.points)
+        colors = np.asarray(pcd.colors) if pcd.has_colors() else None
+
+
+        # ==== 新增：上下翻转点云 ====
+        points[:, 2] = -points[:, 2]  
+        points[:, 0] = -points[:, 0]  
+        # =========================
+
+
+        glb_filename = os.path.join(pc_save_dir, f'{timestamp}_{idx:02d}.glb')
+        success = create_simple_glb_from_pointcloud(points, colors, glb_filename)
+        if not success:
+            logger.warning(f"Failed to save GLB file: {glb_filename}")
+
+        glb_files.append(glb_filename)
+
+    
+    
+    #* save depth predictions video
+    output_filename = f"output_{timestamp}.mp4"
+    output_path = os.path.join(cur_save_dir, output_filename)
+    save_video(color_predictions, output_path, fps=input_fps, quality=5)
+    return output_path, glb_files
+
+
+
+
+
+
+#* gradio creation and initialization
+
+
+css = """
+#video-display-container {
+    max-height: 100vh;
+}
+#video-display-input {
+    max-height: 80vh;
+}
+#video-display-output {
+    max-height: 80vh;
+}
+#download {
+    height: 62px;
+}
+.title {
+    text-align: center;
+}
+.description {
+    text-align: center;
+}
+.gradio-examples {
+    max-height: 400px;
+    overflow-y: auto;
+}
+.gradio-examples .examples-container {
+    display: grid;
+    grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
+    gap: 10px;
+    padding: 10px;
+}
+.gradio-container .gradio-examples .pagination,
+.gradio-container .gradio-examples .pagination button,
+div[data-testid="examples"] .pagination,
+div[data-testid="examples"] .pagination button {
+    font-size: 28px !important;
+    font-weight: bold !important;
+    padding: 15px 20px !important;
+    min-width: 60px !important;
+    height: 60px !important;
+    border-radius: 10px !important;
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+    color: white !important;
+    border: none !important;
+    cursor: pointer !important;
+    margin: 8px !important;
+    display: inline-block !important;
+    box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important;
+    transition: all 0.3s ease !important;
+}
+
+div[data-testid="examples"] .pagination button:not(.active),
+.gradio-container .gradio-examples .pagination button:not(.active) {
+    font-size: 32px !important;
+    font-weight: bold !important;
+    padding: 15px 20px !important;
+    min-width: 60px !important;
+    height: 60px !important;
+    background: linear-gradient(135deg, #8a9cf0 0%, #9a6bb2 100%) !important;
+    opacity: 0.8 !important;
+}
+
+div[data-testid="examples"] .pagination button:hover,
+.gradio-container .gradio-examples .pagination button:hover {
+    background: linear-gradient(135deg, #5a6fd8 0%, #6a4190 100%) !important;
+    transform: translateY(-2px) !important;
+    box-shadow: 0 6px 12px rgba(0,0,0,0.3) !important;
+    opacity: 1 !important;
+}
+
+div[data-testid="examples"] .pagination button.active,
+.gradio-container .gradio-examples .pagination button.active {
+    background: linear-gradient(135deg, #11998e 0%, #38ef7d 100%) !important;
+    box-shadow: 0 4px 8px rgba(17,153,142,0.4) !important;
+    opacity: 1 !important;
+}
+
+button[class*="pagination"],
+button[class*="page"] {
+    font-size: 28px !important;
+    font-weight: bold !important;
+    padding: 15px 20px !important;
+    min-width: 60px !important;
+    height: 60px !important;
+    border-radius: 10px !important;
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+    color: white !important;
+    border: none !important;
+    cursor: pointer !important;
+    margin: 8px !important;
+    box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important;
+    transition: all 0.3s ease !important;
+}
+"""
+
+
+
+head_html = """
+<link rel="icon" type="image/svg+xml" href="data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 100 100'%3E%3Ctext y='.9em' font-size='90'%3E🦾%3C/text%3E%3C/svg%3E">
+<link rel="shortcut icon" type="image/svg+xml" href="data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 100 100'%3E%3Ctext y='.9em' font-size='90'%3E🦾%3C/text%3E%3C/svg%3E">
+<link rel="icon" type="image/png" href="data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 100 100'%3E%3Ctext y='.9em' font-size='90'%3E🦾%3C/text%3E%3C/svg%3E">
+<meta name="viewport" content="width=device-width, initial-scale=1.0">
+"""
+
+
+
+# description = """Official demo for **DKT **."""
+
+# with gr.Blocks(css=css, title="DKT - Diffusion Knows Transparency", favicon_path="favicon.ico") as demo:
+
+
+with gr.Blocks(css=css, title="DKT", head=head_html) as demo:
+    # gr.Markdown(title, elem_classes=["title"])
+    """
+
+    <a title="Website" href="https://stable-x.github.io/StableNormal/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
+        </a>
+        <a title="arXiv" href="https://arxiv.org/abs/2406.16864" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
+        </a>
+        <a title="Social" href="https://x.com/ychngji6" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
+        </a>
+    
+    
+    """
+
+    gr.Markdown(
+        """
+        # Diffusion Knows Transparency: Repurposing Video Diffusion for Transparent Object Depth and Normal Estimation 
+        <p align="center">
+        <a title="Github" href="https://github.com/Daniellli/DKT" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <img src="https://img.shields.io/github/stars/Daniellli/DKT?style=social" alt="badge-github-stars">
+        </a>
+    """
+    )
+    # gr.Markdown(description, elem_classes=["description"])
+    # gr.Markdown("### Video Processing Demo", elem_classes=["description"])
+
+    with gr.Row():
+        with gr.Column():
+            input_video = gr.Video(label="Input Video", elem_id='video-display-input')
+            
+            model_size = gr.Radio(
+                # choices=["1.3B", "14B"],
+                choices=["1.3B"],
+                value="1.3B",
+                label="Model Size"
+            )
+
+
+            with gr.Accordion("Advanced Parameters", open=False):
+                num_inference_steps = gr.Slider(
+                    minimum=1, maximum=50, value=5, step=1,
+                    label="Number of Inference Steps"
+                )
+                overlap = gr.Slider(
+                    minimum=1, maximum=20, value=3, step=1,
+                    label="Overlap"
+                )
+                
+            submit = gr.Button(value="Compute Depth", variant="primary")
+        
+        with gr.Column():
+            output_video = gr.Video(
+                label="Depth Outputs", 
+                elem_id='video-display-output',
+                autoplay=True
+            )
+            vis_video = gr.Video(
+                label="Visualization Video", 
+                visible=False,
+                autoplay=True
+            )
+
+    with gr.Row():
+        gr.Markdown("### 3D Point Cloud Visualization", elem_classes=["title"])
+    
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1):
+            output_point_map0 = gr.Model3D(
+                label="Point Cloud Key Frame 1",
+                clear_color=[1.0, 1.0, 1.0, 1.0],
+                interactive=False,
+            )
+        with gr.Column(scale=1):
+            output_point_map1 = gr.Model3D(
+                label="Point Cloud Key Frame 2",
+                clear_color=[1.0, 1.0, 1.0, 1.0],
+                interactive=False
+            )
+        
+    
+    with gr.Row(equal_height=True):
+        
+        with gr.Column(scale=1):
+            output_point_map2 = gr.Model3D(
+                label="Point Cloud Key Frame 3",
+                clear_color=[1.0, 1.0, 1.0, 1.0],
+                interactive=False
+            )
+        with gr.Column(scale=1):
+            output_point_map3 = gr.Model3D(
+                label="Point Cloud Key Frame 4",
+                clear_color=[1.0, 1.0, 1.0, 1.0],
+                interactive=False
+            )
+
+    def on_submit(video_file, model_size, num_inference_steps, overlap):
+        print('on_submit is calling')
+        logger.info('on_submit is calling')
+
+        if video_file is None:
+            return None, None, None, None, None, None, "Please upload a video file"
+        
+        try:
+
+            output_path, glb_files = process_video(
+                video_file, model_size, num_inference_steps, overlap
+            )
+
+            
+            
+            if output_path is None:
+                return None, None, None, None, None, None, glb_files
+            
+            model3d_outputs = [None] * 4
+            if glb_files:
+                for i, glb_file in enumerate(glb_files[:4]):
+                    if os.path.exists(glb_file):
+                        model3d_outputs[i] = glb_file
+            
+
+
+            return output_path, None, *model3d_outputs
+                
+        except Exception as e:
+            logger.error(e)
+            return None, None, None, None, None, None
+
+    
+    submit.click(
+        on_submit,
+        inputs=[
+            input_video, model_size, num_inference_steps, overlap
+        ],
+        outputs=[
+            output_video, vis_video, output_point_map0, output_point_map1, output_point_map2, output_point_map3
+        ]
+    )
+    
+
+    
+    logger.info(f'there are {len(example_inputs)} demo files')
+    print(f'there are {len(example_inputs)} demo files')
+        
+    examples = gr.Examples(
+        examples=example_inputs, 
+        inputs=[input_video, model_size, num_inference_steps, overlap], 
+        outputs=[
+            output_video, vis_video, 
+            output_point_map0, output_point_map1, output_point_map2, output_point_map3
+        ], 
+        fn=on_submit,
+        examples_per_page=12,
+        cache_examples=False
+    )
+
+
+if __name__ == '__main__':
+    
+    #* main code, model and moge model initialization
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logger.info(f"device = {device}")
+    print(f"device = {device}")
+    
+    load_model_1_3b(device=device)
+    load_moge_model(device=device)
+    # torch.cuda.empty_cache()
+    logger.info('model init done!')
+    print('model init done!')
+
+    demo.queue().launch(share = True)
+    
+    # demo.queue(
+    #     api_open=False,
+    # ).launch()
+
+    
+    # server_name="0.0.0.0", server_port=7860
+
+

dkt/pipelines/{wan_video_new.py → pipeline.py}

@@ -29,7 +29,7 @@ from ..lora import GeneralLoRALoader
 
 from loguru import logger
 
-
+import spaces
 
 class BasePipeline(torch.nn.Module):
 
@@ -222,7 +222,7 @@ class ModelConfig:
                         allow_patterns=allow_file_pattern,
                         ignore_patterns=downloaded_files if downloaded_files else None
                         )
-            
+
             # Let rank 1, 2, ... wait for rank 0
             if use_usp:
                 import torch.distributed as dist
@@ -716,6 +716,323 @@ class WanVideoPipeline(BasePipeline):
 
 
 
+def extract_frames_from_video_file(video_path):
+    try:
+        cap = cv2.VideoCapture(video_path)
+        frames = []
+        
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        if fps <= 0:
+            fps = 15.0
+        
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame_rgb = Image.fromarray(frame_rgb)
+            frames.append(frame_rgb)
+        
+        cap.release()
+        return frames, fps
+    except Exception as e:
+        logger.error(f"Error extracting frames from {video_path}: {str(e)}")
+        return [], 15.0
+
+
+def resize_frame(frame, height, width):
+    frame = np.array(frame)
+    frame = torch.from_numpy(frame).permute(2, 0, 1).unsqueeze(0).float() / 255.0
+    frame = torch.nn.functional.interpolate(frame, (height, width), mode="bicubic", align_corners=False, antialias=True)
+    frame = (frame.squeeze(0).permute(1, 2, 0).clamp(0, 1) * 255).byte().numpy()
+    frame = Image.fromarray(frame)
+    return frame
+
+
+
+from moge.model.v2 import MoGeModel
+from tools.eval_utils import transfer_pred_disp2depth, transfer_pred_disp2depth_v2, colorize_depth_map
+from tools.depth2pcd import depth2pcd
+import cv2, copy
+
+class DKTPipeline:
+    def __init__(self,  ):
+
+        self.main_pipe = self.init_model()
+        
+        self.moge_pipe = self.load_moge_model()
+
+
+
+
+
+    def init_model(self ):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        pipe = WanVideoPipeline.from_pretrained(
+            torch_dtype=torch.bfloat16,
+            device=device,
+            model_configs=[
+                ModelConfig(
+                    model_id="PAI/Wan2.1-Fun-1.3B-Control",
+                    origin_file_pattern="diffusion_pytorch_model*.safetensors",
+                    offload_device="cpu",
+                ),
+                ModelConfig(
+                    model_id="PAI/Wan2.1-Fun-1.3B-Control",
+                    origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth",
+                    offload_device="cpu",
+                ),
+                ModelConfig(
+                    model_id="PAI/Wan2.1-Fun-1.3B-Control",
+                    origin_file_pattern="Wan2.1_VAE.pth",
+                    offload_device="cpu",
+                ),
+                ModelConfig(
+                    model_id="PAI/Wan2.1-Fun-1.3B-Control",
+                    origin_file_pattern="models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",
+                    offload_device="cpu",
+                ),
+            ],
+            training_strategy="origin",
+        )
+        
+        
+        lora_config = ModelConfig(
+            model_id="Daniellesry/DKT-Depth-1-3B",
+            origin_file_pattern="dkt-1-3B.safetensors",
+            offload_device="cpu",
+        )
+        lora_config.download_if_necessary(use_usp=False)
+        
+        pipe.load_lora(pipe.dit, lora_config.path, alpha=1.0)#todo is it work?
+        pipe.enable_vram_management()
+        return pipe
+
+    def load_moge_model(self):
+        device= torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        cached_model_path = 'checkpoints/moge_ckpt/moge-2-vitl-normal/model.pt'
+        if os.path.exists(cached_model_path):
+            logger.info(f"Found cached model at {cached_model_path}, loading from cache...")
+            moge_pipe = MoGeModel.from_pretrained(cached_model_path).to(device)
+        else:
+            logger.info(f"Cache not found at {cached_model_path}, downloading from HuggingFace...")
+            os.makedirs(os.path.dirname(cached_model_path), exist_ok=True)
+            moge_pipe = MoGeModel.from_pretrained('Ruicheng/moge-2-vitl-normal', cache_dir=os.path.dirname(cached_model_path)).to(device)
+        
+        return moge_pipe
+
+
+
+    @spaces.GPU(duration=120)
+    @torch.inference_mode()
+    def __call__(self, video_file, prompt='depth', \
+        negative_prompt='', height=480, width=832, \
+            num_inference_steps=5, window_size=21, \
+                overlap=3, vis_pc = False, return_rgb = False):
+        
+        
+        origin_frames, input_fps = extract_frames_from_video_file(video_file)
+
+        frame_length = len(origin_frames)
+
+        original_width, original_height = origin_frames[0].size
+
+        ROTATE = False 
+        if original_width <  original_height:#* ensure the width is the longer side
+            ROTATE = True
+            origin_frames = [x.transpose(Image.ROTATE_90) for x in origin_frames]
+            tmp = original_width
+            original_width = original_height
+            original_height = tmp
+
+
+        frames = [resize_frame(frame, height, width) for frame in origin_frames]
+
+
+        if (frame_length - 1) % 4 != 0:
+            new_len = ((frame_length - 1) // 4 + 1) * 4 + 1
+            frames = frames + [copy.deepcopy(frames[-1]) for _ in range(new_len - frame_length)]
+
+        
+
+        
+        video, vae_outs = self.main_pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            control_video=frames,
+            height=height,
+            width=width,
+            num_frames=len(frames),
+            seed=1,
+            tiled=False,
+            num_inference_steps=num_inference_steps,
+            sliding_window_size=window_size,
+            sliding_window_stride=window_size - overlap,
+            cfg_scale=1.0,
+        )
+        torch.cuda.empty_cache()
+
+        processed_video = video[:frame_length]
+        processed_video = [resize_frame(frame, original_height, original_width) for frame in processed_video]
+
+        if ROTATE:
+            processed_video = [x.transpose(Image.ROTATE_270) for x in processed_video]
+            origin_frames = [x.transpose(Image.ROTATE_270) for x in origin_frames]
+
+        color_predictions = []
+        if prompt == 'depth':
+            prediced_depth_map_np = [np.array(item).astype(np.float32).mean(-1)  for item in processed_video]
+            prediced_depth_map_np = np.stack(prediced_depth_map_np)
+            prediced_depth_map_np = prediced_depth_map_np / 255.0
+            
+            __min = prediced_depth_map_np.min()
+            __max = prediced_depth_map_np.max()
+
+            prediced_depth_map_np_normalized = (prediced_depth_map_np - __min) / (__max - __min)
+            color_predictions = [colorize_depth_map(item) for item in prediced_depth_map_np_normalized]
+        else:
+            color_predictions = processed_video
+
+        return_dict = {}
+        
+        return_dict['depth_map'] = prediced_depth_map_np
+        return_dict['colored_depth_map'] = color_predictions
+
+
+
+        if vis_pc and prompt == 'depth':
+            vis_pc_num = 4
+            indices = np.linspace(0, frame_length-1, vis_pc_num)
+            indices = np.round(indices).astype(np.int32)
+            return_dict['point_clouds'] = self.prediction2pc(prediced_depth_map_np, origin_frames, indices)
+
+        if return_rgb:
+            return_dict['rgb_frames'] = origin_frames
+        return return_dict
+
+        
+        
+        
+            
+    def prediction2pc(self, prediction_depth_map, RGB_frames, indices, return_pcd = True,nb_neighbors = 20, std_ratio = 3.0):
+        resize_W,resize_H = RGB_frames[0].size
+        pcds = []
+        moge_device = self.moge_pipe.device if self.moge_pipe is not None else torch.device("cuda:0")
+        
+        for idx in tqdm(indices):
+            orgin_rgb_frame = RGB_frames[idx]
+            predicted_depth = prediction_depth_map[idx]
+
+            # Read the input image and convert to tensor (3, H, W) with RGB values normalized to [0, 1]
+            input_image_np = np.array(orgin_rgb_frame)  # Convert PIL Image to numpy array
+            input_image = torch.tensor(input_image_np / 255, dtype=torch.float32, device=moge_device).permute(2, 0, 1) 
+            output = self.moge_pipe.infer(input_image)
+
+            #* "dict_keys(['points', 'intrinsics', 'depth', 'mask', 'normal'])"
+            moge_intrinsics = output['intrinsics'].cpu().numpy()
+            moge_mask = output['mask'].cpu().numpy()
+            moge_depth = output['depth'].cpu().numpy()
+        
+    
+            metric_depth = transfer_pred_disp2depth(predicted_depth, moge_depth, moge_mask)
+            
+            moge_intrinsics[0, 0] *= resize_W 
+            moge_intrinsics[1, 1] *= resize_H
+            moge_intrinsics[0, 2] *= resize_W
+            moge_intrinsics[1, 2] *= resize_H
+            
+            pcd = depth2pcd(metric_depth, moge_intrinsics, color=input_image_np, input_mask=moge_mask, ret_pcd=return_pcd)
+
+            if return_pcd:
+                #* [15,50], [2,3] 
+                cl, ind = pcd.remove_statistical_outlier(nb_neighbors=nb_neighbors, std_ratio=std_ratio)
+                pcd = pcd.select_by_index(ind)
+                #todo downsample
+            
+            pcds.append(pcd)
+        
+        return pcds
+
+
+
+
+    @spaces.GPU()
+    @torch.inference_mode()
+    def moge_infer(self, input_image):
+        return self.moge_pipe.infer(input_image)
+        
+
+
+    def prediction2pc_v2(self, prediction_depth_map, RGB_frames, indices, return_pcd = True,nb_neighbors = 20, std_ratio = 3.0):
+        """
+            call MoGe once 
+        """
+        resize_W,resize_H = RGB_frames[0].size
+        pcds = []
+        moge_device = self.moge_pipe.device if self.moge_pipe is not None else torch.device("cuda:0")
+        
+        for iidx, idx in enumerate(tqdm(indices)):
+
+            orgin_rgb_frame = RGB_frames[idx]
+            predicted_depth = prediction_depth_map[idx]
+            input_image_np = np.array(orgin_rgb_frame)  # Convert PIL Image to numpy array
+
+
+            if iidx == 0:
+                # Read the input image and convert to tensor (3, H, W) with RGB values normalized to [0, 1]
+                input_image = torch.tensor(input_image_np / 255, dtype=torch.float32, device=moge_device).permute(2, 0, 1) 
+                output = self.moge_infer(input_image)
+
+                #* "dict_keys(['points', 'intrinsics', 'depth', 'mask', 'normal'])"
+                moge_intrinsics = output['intrinsics'].cpu().numpy()
+                moge_mask = output['mask'].cpu().numpy()
+                moge_depth = output['depth'].cpu().numpy()
+
+                metric_depth, scale, shift = transfer_pred_disp2depth(predicted_depth, moge_depth, moge_mask, return_scale_shift=True)
+
+                moge_intrinsics[0, 0] *= resize_W 
+                moge_intrinsics[1, 1] *= resize_H
+                moge_intrinsics[0, 2] *= resize_W
+                moge_intrinsics[1, 2] *= resize_H
+            else:
+                metric_depth = transfer_pred_disp2depth_v2(predicted_depth, scale, shift)
+
+            
+            pcd = depth2pcd(metric_depth, moge_intrinsics, color=input_image_np, input_mask=moge_mask, ret_pcd=return_pcd)
+
+            if return_pcd:
+                #* [15,50], [2,3] 
+                cl, ind = pcd.remove_statistical_outlier(nb_neighbors=nb_neighbors, std_ratio=std_ratio)
+                pcd = pcd.select_by_index(ind)
+                #todo downsample
+            
+            pcds.append(pcd)
+        
+        return pcds
+
+
+    
+
+        
+
+        
+
+        
+
+
+
+
+
+        
+        
+
+
+    
+        
+
+        
 
 
 
@@ -1480,7 +1797,4 @@ def model_fn_wan_video(
 
     #* unpatchify, from [1, ( (F-1)/4 * H/16 *  W/16), 64] to [1, 16, (F-1)/4, H/8, W/8]
     x = dit.unpatchify(x, (f, h, w))
-    return x
-
-
-
+    return x