Upload 4 files

app.py

@@ -0,0 +1,651 @@
+try:
+    import spaces
+    GPU = spaces.GPU
+    print("spaces GPU is available")
+except ImportError:
+    def GPU(func):
+        return func
+
+import os
+import subprocess
+
+# def install_cuda_toolkit():
+#     # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
+#     CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run"
+#     CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
+#     subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
+#     subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
+#     subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
+    
+#     os.environ["CUDA_HOME"] = "/usr/local/cuda"
+#     os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
+#     os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
+#         os.environ["CUDA_HOME"],
+#         "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
+#     )
+#     # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
+#     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
+
+#     print("Successfully installed CUDA toolkit at: ", os.environ["CUDA_HOME"])
+
+#     subprocess.call('rm /usr/bin/gcc', shell=True)
+#     subprocess.call('rm /usr/bin/g++', shell=True)
+#     subprocess.call('rm /usr/local/cuda/bin/gcc', shell=True)
+#     subprocess.call('rm /usr/local/cuda/bin/g++', shell=True)
+
+#     subprocess.call('ln -s /usr/bin/gcc-11 /usr/bin/gcc', shell=True)
+#     subprocess.call('ln -s /usr/bin/g++-11 /usr/bin/g++', shell=True)
+
+#     subprocess.call('ln -s /usr/bin/gcc-11 /usr/local/cuda/bin/gcc', shell=True)
+#     subprocess.call('ln -s /usr/bin/g++-11 /usr/local/cuda/bin/g++', shell=True)
+
+#     subprocess.call('gcc --version', shell=True)
+#     subprocess.call('g++ --version', shell=True)
+
+# install_cuda_toolkit()
+
+# subprocess.run('pip install git+https://github.com/nerfstudio-project/gsplat.git@32f2a54d21c7ecb135320bb02b136b7407ae5712 --no-build-isolation --use-pep517', env={'CUDA_HOME': "/usr/local/cuda", "TORCH_CUDA_ARCH_LIST": "8.0;8.6"}, shell=True)
+
+from flask import Flask, jsonify, request, send_file, render_template
+import base64
+import io
+from PIL import Image
+import torch
+import numpy as np
+import os
+import argparse
+import imageio
+import json
+
+import time
+import threading
+
+from concurrency_manager import ConcurrencyManager
+
+from huggingface_hub import hf_hub_download
+
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+import imageio
+
+from models import *
+from utils import *
+
+from transformers import T5TokenizerFast, UMT5EncoderModel
+
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+class MyFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        return torch.argmin(
+            (timestep - schedule_timesteps.to(timestep.device)).abs(), dim=0).item()
+
+class GenerationSystem(nn.Module):
+    def __init__(self, ckpt_path=None, device="cuda:0", offload_t5=False, offload_vae=False):
+        super().__init__()
+        self.device = device
+        self.offload_t5 = offload_t5
+        self.offload_vae = offload_vae
+
+        self.latent_dim = 48
+        self.temporal_downsample_factor = 4
+        self.spatial_downsample_factor = 16
+
+        self.feat_dim = 1024
+
+        self.latent_patch_size = 2
+
+        self.denoising_steps = [0, 250, 500, 750]
+
+        model_id = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+
+        self.vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float).eval()
+
+        from models.autoencoder_kl_wan import WanCausalConv3d
+        with torch.no_grad():
+            for name, module in self.vae.named_modules():
+                if isinstance(module, WanCausalConv3d):
+                    time_pad = module._padding[4]
+                    module.padding = (0, module._padding[2], module._padding[0])
+                    module._padding = (0, 0, 0, 0, 0, 0)
+                    module.weight = torch.nn.Parameter(module.weight[:, :, time_pad:].clone())
+
+        self.vae.requires_grad_(False)
+
+        self.register_buffer('latents_mean', torch.tensor(self.vae.config.latents_mean).float().view(1, self.vae.config.z_dim, 1, 1, 1).to(self.device))
+        self.register_buffer('latents_std', torch.tensor(self.vae.config.latents_std).float().view(1, self.vae.config.z_dim, 1, 1, 1).to(self.device))
+
+        self.latent_scale_fn = lambda x: (x - self.latents_mean) / self.latents_std
+        self.latent_unscale_fn = lambda x: x * self.latents_std + self.latents_mean
+
+        self.tokenizer = T5TokenizerFast.from_pretrained(model_id, subfolder="tokenizer")
+
+        self.text_encoder = UMT5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float32).eval().requires_grad_(False).to(self.device if not self.offload_t5 else "cpu")
+
+        self.transformer = WanTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float32).train().requires_grad_(False)
+        
+        self.transformer.patch_embedding.weight = nn.Parameter(F.pad(self.transformer.patch_embedding.weight, (0, 0, 0, 0, 0, 0, 0, 6 + self.latent_dim)))
+        # self.transformer.rope.freqs_f[:] = self.transformer.rope.freqs_f[:1]
+
+        weight = self.transformer.proj_out.weight.reshape(self.latent_patch_size ** 2, self.latent_dim, self.transformer.proj_out.weight.shape[1])
+        bias = self.transformer.proj_out.bias.reshape(self.latent_patch_size ** 2, self.latent_dim)
+
+        extra_weight = torch.randn(self.latent_patch_size ** 2, self.feat_dim, self.transformer.proj_out.weight.shape[1]) * 0.02
+        extra_bias = torch.zeros(self.latent_patch_size ** 2, self.feat_dim)
+ 
+        self.transformer.proj_out.weight = nn.Parameter(torch.cat([weight, extra_weight], dim=1).flatten(0, 1).detach().clone())
+        self.transformer.proj_out.bias = nn.Parameter(torch.cat([bias, extra_bias], dim=1).flatten(0, 1).detach().clone())
+
+        self.recon_decoder = WANDecoderPixelAligned3DGSReconstructionModel(self.vae, self.feat_dim, use_render_checkpointing=True, use_network_checkpointing=False).train().requires_grad_(False).to(self.device)
+
+        self.scheduler = MyFlowMatchEulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler", shift=3)
+
+        self.register_buffer('timesteps', self.scheduler.timesteps.clone().to(self.device))
+
+        self.transformer.disable_gradient_checkpointing()
+        self.transformer.gradient_checkpointing = False
+
+        self.add_feedback_for_transformer()
+
+        if ckpt_path is not None:
+            state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+            self.transformer.load_state_dict(state_dict["transformer"])
+            self.recon_decoder.load_state_dict(state_dict["recon_decoder"])
+            print(f"Loaded {ckpt_path}.")
+
+        from quant import FluxFp8GeMMProcessor
+
+        FluxFp8GeMMProcessor(self.transformer)
+
+        del self.vae.post_quant_conv, self.vae.decoder
+        self.vae.to(self.device if not self.offload_vae else "cpu")
+
+        self.transformer.to(self.device)
+
+    def add_feedback_for_transformer(self):
+        self.use_feedback = True
+        self.transformer.patch_embedding.weight = nn.Parameter(F.pad(self.transformer.patch_embedding.weight, (0, 0, 0, 0, 0, 0, 0, self.feat_dim + self.latent_dim)))
+    
+    def encode_text(self, texts):
+        max_sequence_length = 512
+
+        text_inputs = self.tokenizer(
+            texts,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+        if getattr(self, "offload_t5", False):
+            text_input_ids = text_inputs.input_ids.to("cpu")
+            mask = text_inputs.attention_mask.to("cpu")
+        else:
+            text_input_ids = text_inputs.input_ids.to(self.device)
+            mask = text_inputs.attention_mask.to(self.device)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+
+        if getattr(self, "offload_t5", False):
+            with torch.no_grad():
+                text_embeds = self.text_encoder(text_input_ids, mask).last_hidden_state.to(self.device)
+        else:
+            text_embeds = self.text_encoder(text_input_ids, mask).last_hidden_state
+        text_embeds = [u[:v] for u, v in zip(text_embeds, seq_lens)]
+        text_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in text_embeds], dim=0
+        )
+        return text_embeds.float()
+
+    def forward_generator(self, noisy_latents, raymaps, condition_latents, t, text_embeds, cameras, render_cameras, image_height, image_width, need_3d_mode=True):
+
+        out = self.transformer(
+            hidden_states=torch.cat([noisy_latents, raymaps, condition_latents], dim=1),
+            timestep=t,
+            encoder_hidden_states=text_embeds,
+            return_dict=False,
+        )[0]
+
+        v_pred, feats = out.split([self.latent_dim, self.feat_dim], dim=1)
+               
+        sigma = torch.stack([self.scheduler.sigmas[self.scheduler.index_for_timestep(_t)] for _t in t.unbind(0)], dim=0).to(self.device)
+        latents_pred_2d = noisy_latents - sigma * v_pred
+
+        if need_3d_mode:
+            scene_params = self.recon_decoder(
+                                einops.rearrange(feats, 'B C T H W -> (B T) C H W').unsqueeze(2), 
+                                einops.rearrange(self.latent_unscale_fn(latents_pred_2d.detach()), 'B C T H W -> (B T) C H W').unsqueeze(2), 
+                                cameras
+                            ).flatten(1, -2)
+
+            images_pred, _ = self.recon_decoder.render(scene_params.unbind(0), render_cameras, image_height, image_width, bg_mode="white")
+
+            latents_pred_3d = einops.rearrange(self.latent_scale_fn(self.vae.encode(
+                            einops.rearrange(images_pred, 'B T C H W -> (B T) C H W', T=images_pred.shape[1]).unsqueeze(2).to(self.device if not self.offload_vae else "cpu").float()
+                        ).latent_dist.sample().to(self.device)).squeeze(2), '(B T) C H W -> B C T H W', T=images_pred.shape[1]).to(noisy_latents.dtype)
+
+        return {
+            '2d': latents_pred_2d,
+            '3d': latents_pred_3d if need_3d_mode else None,
+            'rgb_3d': images_pred if need_3d_mode else None,
+            'scene': scene_params if need_3d_mode else None,
+            'feat': feats
+        }
+
+    @torch.no_grad()
+    @torch.amp.autocast(dtype=torch.bfloat16, device_type="cuda")
+    def generate(self, cameras, n_frame, image=None, text="", image_index=0, image_height=480, image_width=704, video_output_path=None):            
+        with torch.no_grad():
+            batch_size = 1
+            
+            cameras = cameras.to(self.device).unsqueeze(0)
+
+            if cameras.shape[1] != n_frame:
+                render_cameras = cameras.clone()
+                cameras = sample_from_dense_cameras(cameras.squeeze(0), torch.linspace(0, 1, n_frame, device=self.device)).unsqueeze(0)
+            else:
+                render_cameras = cameras
+            
+            cameras, ref_w2c, T_norm = normalize_cameras(cameras, return_meta=True, n_frame=None)
+
+            render_cameras = normalize_cameras(render_cameras, ref_w2c=ref_w2c, T_norm=T_norm, n_frame=None)
+
+            text = "[Static] " + text
+
+            text_embeds = self.encode_text([text])
+            # neg_text_embeds = self.encode_text([""]).repeat(batch_size, 1, 1)
+
+            masks = torch.zeros(batch_size, n_frame, device=self.device)
+
+            condition_latents = torch.zeros(batch_size, self.latent_dim, n_frame, image_height // self.spatial_downsample_factor, image_width // self.spatial_downsample_factor, device=self.device)
+
+            if image is not None:
+                image = image.to(self.device)
+
+                latent = self.latent_scale_fn(self.vae.encode(
+                        image.unsqueeze(0).unsqueeze(2).to(self.device if not self.offload_vae else "cpu").float()
+                    ).latent_dist.sample().to(self.device)).squeeze(2)
+
+                masks[:, image_index] = 1
+                condition_latents[:, :, image_index] = latent
+
+            raymaps = create_raymaps(cameras, image_height // self.spatial_downsample_factor, image_width // self.spatial_downsample_factor)
+            raymaps = einops.rearrange(raymaps, 'B T H W C -> B C T H W', T=n_frame)
+            
+            noise = torch.randn(batch_size, self.latent_dim, n_frame, image_height // self.spatial_downsample_factor, image_width // self.spatial_downsample_factor, device=self.device)
+
+            noisy_latents = noise 
+
+            torch.cuda.empty_cache()
+
+            if self.use_feedback:
+                prev_latents_pred = torch.zeros(batch_size, self.latent_dim, n_frame, image_height // self.spatial_downsample_factor, image_width // self.spatial_downsample_factor, device=self.device)
+
+                prev_feats = torch.zeros(batch_size, self.feat_dim, n_frame, image_height // self.spatial_downsample_factor, image_width // self.spatial_downsample_factor, device=self.device)
+
+            for i in range(len(self.denoising_steps)):
+                t_ids = torch.full((noisy_latents.shape[0],), self.denoising_steps[i], device=self.device)
+
+                t = self.timesteps[t_ids]
+
+                if self.use_feedback:
+                    _condition_latents = torch.cat([condition_latents, prev_feats, prev_latents_pred], dim=1)
+                else:
+                    _condition_latents = condition_latents
+
+                if i < len(self.denoising_steps) - 1:
+                    out = self.forward_generator(noisy_latents, raymaps, _condition_latents, t, text_embeds, cameras, cameras, image_height, image_width, need_3d_mode=True)
+
+                    latents_pred = out["3d"]
+
+                    if self.use_feedback:
+                        prev_latents_pred = latents_pred
+                        prev_feats = out['feat']
+                   
+                    noisy_latents = self.scheduler.scale_noise(latents_pred, self.timesteps[torch.full((noisy_latents.shape[0],), self.denoising_steps[i + 1], device=self.device)], torch.randn_like(noise))
+                    
+                else:
+                    out = self.transformer(
+                        hidden_states=torch.cat([noisy_latents, raymaps, _condition_latents], dim=1),
+                        timestep=t,
+                        encoder_hidden_states=text_embeds,
+                        return_dict=False,
+                    )[0]
+
+                    v_pred, feats = out.split([self.latent_dim, self.feat_dim], dim=1)
+                        
+                    sigma = torch.stack([self.scheduler.sigmas[self.scheduler.index_for_timestep(_t)] for _t in t.unbind(0)], dim=0).to(self.device)
+                    latents_pred = noisy_latents - sigma * v_pred
+
+                    scene_params = self.recon_decoder(
+                                        einops.rearrange(feats, 'B C T H W -> (B T) C H W').unsqueeze(2), 
+                                        einops.rearrange(self.latent_unscale_fn(latents_pred.detach()), 'B C T H W -> (B T) C H W').unsqueeze(2), 
+                                        cameras
+                                    ).flatten(1, -2)
+
+            if video_output_path is not None:
+                interpolated_images_pred, _ = self.recon_decoder.render(scene_params.unbind(0), render_cameras, image_height, image_width, bg_mode="white")
+
+                interpolated_images_pred = einops.rearrange(interpolated_images_pred[0].clamp(-1, 1).add(1).div(2), 'T C H W -> T H W C')
+
+                interpolated_images_pred = [torch.cat([img], dim=1).detach().cpu().mul(255).numpy().astype(np.uint8) for i, img in enumerate(interpolated_images_pred.unbind(0))]
+
+                imageio.mimwrite(video_output_path, interpolated_images_pred, fps=15, quality=8, macro_block_size=1) 
+
+        scene_params = scene_params[0]    
+
+        scene_params = scene_params.detach().cpu()
+
+        return scene_params, ref_w2c, T_norm
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--port', type=int, default=7860)
+    parser.add_argument("--ckpt", default=None)
+    parser.add_argument("--gpu", type=int, default=0)
+    parser.add_argument("--cache_dir", type=str, default="./tmpfiles")
+    parser.add_argument("--offload_t5", type=bool, default=False)
+    parser.add_argument("--max_concurrent", type=int, default=1, help="Maximum concurrent generation tasks")
+    args, _ = parser.parse_known_args()
+
+    # Ensure model.ckpt exists, download if not present
+    if args.ckpt is None:
+        from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
+        ckpt_path = os.path.join(HUGGINGFACE_HUB_CACHE, "models--imlixinyang--FlashWorld", "snapshots", "6a8e88c6f88678ac098e4c82675f0aee555d6e5d", "model.ckpt")
+        if not os.path.exists(ckpt_path):
+            hf_hub_download(repo_id="imlixinyang/FlashWorld", filename="model.ckpt", local_dir_use_symlinks=False)
+    else:
+        ckpt_path = args.ckpt
+
+    app = Flask(__name__)
+    
+    # 初始化GenerationSystem
+    device = f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu"
+    generation_system = GenerationSystem(ckpt_path=ckpt_path, device=device)
+    
+    # 初始化并发管理器
+    concurrency_manager = ConcurrencyManager(max_concurrent=args.max_concurrent)
+
+    @app.after_request
+    def after_request(response):
+        response.headers.add('Access-Control-Allow-Origin', '*')
+        response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
+        response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
+        return response
+    
+    @GPU
+    def generate_wrapper(cameras, n_frame, image, text_prompt, image_index, image_height, image_width, video_output_path=None):
+        """生成函数的包装器，用于并发控制"""
+        return generation_system.generate(cameras, n_frame, image, text_prompt, image_index, image_height, image_width, video_output_path)
+
+    def job_generate(file_id, cache_dir, payload):
+        """工作线程执行的生成任务：负责生成并落盘，返回可下载信息"""
+        # 解包参数
+        cameras = payload["cameras"]
+        n_frame = payload["n_frame"]
+        image = payload["image"]
+        text_prompt = payload["text_prompt"]
+        image_index = payload["image_index"]
+        image_height = payload["image_height"]
+        image_width = payload["image_width"]
+        data = payload["raw_request"]
+
+        # 执行生成
+        scene_params, ref_w2c, T_norm = generation_system.generate(
+            cameras, n_frame, image, text_prompt, image_index, image_height, image_width, video_output_path=None
+        )
+
+        # 保存请求元数据
+        with open(os.path.join(cache_dir, f'{file_id}.json'), 'w') as f:
+            json.dump(data, f)
+
+        # 导出PLY文件
+        splat_path = os.path.join(cache_dir, f'{file_id}.ply')
+        export_ply_for_gaussians(splat_path, scene_params, opacity_threshold=0.001, T_norm=T_norm)
+
+        file_size = os.path.getsize(splat_path) if os.path.exists(splat_path) else 0
+
+        return {
+            'file_id': file_id,
+            'file_path': splat_path,
+            'file_size': file_size,
+            'download_url': f'/download/{file_id}'
+        }
+    
+    @app.route('/generate', methods=['POST', 'OPTIONS'])
+    def generate():
+        # Handle preflight request
+        if request.method == 'OPTIONS':
+            return jsonify({'status': 'ok'})
+        
+        try:
+            data = request.get_json(force=True)
+            
+            image_prompt = data.get('image_prompt', None)
+            text_prompt = data.get('text_prompt', "")
+            cameras = data.get('cameras')
+            resolution = data.get('resolution')
+            image_index = data.get('image_index', 0)
+
+            n_frame, image_height, image_width = resolution
+
+            if not image_prompt and text_prompt == "":
+                return jsonify({'error': 'No Prompts provided'}), 400
+
+            # 处理图像
+            if image_prompt:
+                # image_prompt可以是路径和base64
+                if os.path.exists(image_prompt):
+                    image_prompt = Image.open(image_prompt)
+                else:
+                    # image_prompt 可能是 "data:image/png;base64,...."
+                    if ',' in image_prompt:
+                        image_prompt = image_prompt.split(',', 1)[1]
+                    
+                    try:
+                        image_bytes = base64.b64decode(image_prompt)
+                        image_prompt = Image.open(io.BytesIO(image_bytes))
+                    except Exception as img_e:
+                        return jsonify({'error': f'Image decode error: {str(img_e)}'}), 400
+
+                image = image_prompt.convert('RGB')
+
+                w, h = image.size
+
+                # center crop
+                if image_height / h > image_width / w:
+                    scale = image_height / h
+                else:
+                    scale = image_width / w
+                    
+                new_h = int(image_height / scale)
+                new_w = int(image_width / scale)
+
+                image = image.crop(((w - new_w) // 2, (h - new_h) // 2, 
+                                    new_w + (w - new_w) // 2, new_h + (h - new_h) // 2)).resize((image_width, image_height))
+
+                for camera in cameras:
+                    camera['fx'] = camera['fx'] * scale 
+                    camera['fy'] = camera['fy'] * scale 
+                    camera['cx'] = (camera['cx'] - (w - new_w) // 2) * scale
+                    camera['cy'] = (camera['cy'] - (h - new_h) // 2) * scale
+
+                image = torch.from_numpy(np.array(image)).float().permute(2, 0, 1) / 255.0 * 2 - 1
+            else:
+                image = None
+
+            cameras = torch.stack([
+                torch.from_numpy(np.array([camera['quaternion'][0], camera['quaternion'][1], camera['quaternion'][2], camera['quaternion'][3], camera['position'][0], camera['position'][1], camera['position'][2], camera['fx'] / image_width, camera['fy'] / image_height, camera['cx'] / image_width, camera['cy'] / image_height], dtype=np.float32))
+                for camera in cameras
+            ], dim=0)
+
+            file_id = str(int(time.time() * 1000))
+
+            # 组装任务参数，推迟执行与落盘到工作线程中
+            payload = {
+                'cameras': cameras,
+                'n_frame': n_frame,
+                'image': image,
+                'text_prompt': text_prompt,
+                'image_index': image_index,
+                'image_height': image_height,
+                'image_width': image_width,
+                'raw_request': data,
+            }
+
+            # 提交任务到并发管理器（异步）
+            task_id = concurrency_manager.submit_task(
+                job_generate, file_id, args.cache_dir, payload
+            )
+
+            # 提交后立即返回队列信息
+            queue_status = concurrency_manager.get_queue_status()
+            queued_tasks = queue_status.get('queued_tasks', [])
+            try:
+                queue_position = queued_tasks.index(task_id) + 1
+            except ValueError:
+                # 如果任务已被工作线程立即领取，则认为已开始执行，位置为 0
+                queue_position = 0
+
+            return jsonify({
+                'success': True,
+                'task_id': task_id,
+                'file_id': file_id,
+                'queue': {
+                    'queued_count': queue_status.get('queued_count', 0),
+                    'running_count': queue_status.get('running_count', 0),
+                    'position': queue_position
+                }
+            }), 202
+                
+        except Exception as e:
+            return jsonify({'error': f'Server error: {str(e)}'}), 500
+
+    @app.route('/download/<file_id>', methods=['GET'])
+    def download_file(file_id):
+        """下载生成的PLY文件"""
+        file_path = os.path.join(args.cache_dir, f'{file_id}.ply')
+        
+        if not os.path.exists(file_path):
+            return jsonify({'error': 'File not found'}), 404
+        
+        return send_file(file_path, as_attachment=True, download_name=f'{file_id}.ply')
+
+    @app.route('/delete/<file_id>', methods=['DELETE', 'POST', 'OPTIONS'])
+    def delete_file_endpoint(file_id):
+        """删除生成的文件及其元数据（由前端在下载完成后调用）"""
+        # CORS preflight
+        if request.method == 'OPTIONS':
+            return jsonify({'status': 'ok'})
+
+        try:
+            ply_path = os.path.join(args.cache_dir, f'{file_id}.ply')
+            json_path = os.path.join(args.cache_dir, f'{file_id}.json')
+            deleted = []
+            for path in [ply_path, json_path]:
+                if os.path.exists(path):
+                    os.remove(path)
+                    deleted.append(os.path.basename(path))
+            return jsonify({'success': True, 'deleted': deleted})
+        except Exception as e:
+            return jsonify({'success': False, 'error': str(e)}), 500
+    
+    @app.route('/status', methods=['GET'])
+    def get_status():
+        """获取系统状态和队列信息"""
+        try:
+            queue_status = concurrency_manager.get_queue_status()
+            return jsonify({
+                'success': True,
+                'status': queue_status,
+                'timestamp': time.time()
+            })
+        except Exception as e:
+            return jsonify({'error': f'Failed to get status: {str(e)}'}), 500
+    
+    @app.route('/task/<task_id>', methods=['GET'])
+    def get_task_status(task_id):
+        """获取特定任务的状态（包含排队位置和完成后的文件信息）"""
+        try:
+            task = concurrency_manager.get_task_status(task_id)
+            if not task:
+                return jsonify({'error': 'Task not found'}), 404
+
+            queue_status = concurrency_manager.get_queue_status()
+            queued_tasks = queue_status.get('queued_tasks', [])
+            try:
+                queue_position = queued_tasks.index(task_id) + 1
+            except ValueError:
+                queue_position = 0
+
+            resp = {
+                'success': True,
+                'task_id': task_id,
+                'status': task.status.value,
+                'created_at': task.created_at,
+                'started_at': task.started_at,
+                'completed_at': task.completed_at,
+                'error': task.error,
+                'queue': {
+                    'queued_count': queue_status.get('queued_count', 0),
+                    'running_count': queue_status.get('running_count', 0),
+                    'position': queue_position
+                }
+            }
+
+            if task.status.value == 'completed' and isinstance(task.result, dict):
+                resp.update({
+                    'file_id': task.result.get('file_id'),
+                    'file_path': task.result.get('file_path'),
+                    'file_size': task.result.get('file_size'),
+                    'download_url': task.result.get('download_url'),
+                    'generation_time': (task.completed_at - task.started_at)
+                })
+
+                # 更新task状态
+
+            return jsonify(resp)
+        except Exception as e:
+            return jsonify({'error': f'Failed to get task status: {str(e)}'}), 500
+
+    @app.route("/")
+    def index():
+        return send_file("index.html")
+
+    os.makedirs(args.cache_dir, exist_ok=True)
+
+    # 后台定时清理：删除超过30分钟未访问/修改的缓存文件
+    def cleanup_worker(cache_dir: str, max_age_seconds: int = 1800, interval_seconds: int = 300):
+        while True:
+            try:
+                now = time.time()
+                for name in os.listdir(cache_dir):
+                    # 只清理与任务相关的 .ply/.json 文件
+                    if not (name.endswith('.ply') or name.endswith('.json')):
+                        continue
+                    path = os.path.join(cache_dir, name)
+                    try:
+                        mtime = os.path.getmtime(path)
+                        if now - mtime > max_age_seconds:
+                            os.remove(path)
+                    except FileNotFoundError:
+                        pass
+                    except Exception:
+                        # 忽略单个文件的异常，继续清理
+                        pass
+            except Exception:
+                # 防止线程因异常退出
+                pass
+            time.sleep(interval_seconds)
+
+    cleaner_thread = threading.Thread(target=cleanup_worker, args=(args.cache_dir,), daemon=True)
+    cleaner_thread.start()
+
+    app.run(host='0.0.0.0', port=args.port)

quant.py

@@ -0,0 +1,195 @@
+import gc
+from typing import Tuple
+import copy
+import torch
+import tqdm
+
+
+def cleanup_memory():
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+def per_tensor_quantize(tensor: torch.Tensor) -> Tuple[torch.Tensor, float]:
+    """Quantize a tensor using per-tensor static scaling factor.
+    Args:
+        tensor: The input tensor.
+    """
+    finfo = torch.finfo(torch.float8_e4m3fn)
+    # Calculate the scale as dtype max divided by absmax.
+    # Since .abs() creates a new tensor, we use aminmax to get
+    # the min and max first and then calculate the absmax.
+    if tensor.numel() == 0:
+        # Deal with empty tensors (triggered by empty MoE experts)
+        min_val, max_val = (
+            torch.tensor(-16.0, dtype=tensor.dtype),
+            torch.tensor(16.0, dtype=tensor.dtype),
+        )
+    else:
+        min_val, max_val = tensor.aminmax()
+    amax = torch.maximum(min_val.abs(), max_val.abs())
+    scale = finfo.max / amax.clamp(min=1e-12)
+    # scale and clamp the tensor to bring it to
+    # the representative range of float8 data type
+    # (as default cast is unsaturated)
+    qweight = (tensor * scale).clamp(min=finfo.min, max=finfo.max)
+    # Return both float8 data and the inverse scale (as float),
+    # as both required as inputs to torch._scaled_mm
+    qweight = qweight.to(torch.float8_e4m3fn)
+    scale = scale.float().reciprocal()
+    return qweight, scale
+
+
+def static_per_tensor_quantize(tensor: torch.Tensor, inv_scale: float) -> torch.Tensor:
+    """Quantizes a floating-point tensor to FP8 (E4M3 format) using static scaling.
+    
+    Performs uniform quantization of the input tensor by:
+    1. Scaling the tensor values using the provided inverse scale factor
+    2. Clamping values to the representable range of FP8 E4M3 format
+    3. Converting to FP8 data type
+    
+    Args:
+        tensor (torch.Tensor): Input tensor to be quantized (any floating-point dtype)
+        inv_scale (float): Inverse of the quantization scale factor (1/scale)
+                         (Must be pre-calculated based on tensor statistics)
+    
+    Returns:
+        torch.Tensor: Quantized tensor in torch.float8_e4m3fn format
+    
+    Note:
+        - Uses the E4M3 format (4 exponent bits, 3 mantissa bits, no infinity/nan)
+        - This is a static quantization (scale factor must be pre-determined)
+        - For dynamic quantization, see per_tensor_quantize()
+    """
+    finfo = torch.finfo(torch.float8_e4m3fn)
+    qweight = (tensor / inv_scale).clamp(min=finfo.min, max=finfo.max)
+    return qweight.to(torch.float8_e4m3fn)
+
+
+def fp8_gemm(A, A_scale, B, B_scale, bias, out_dtype, native_fp8_support=False):
+    """Performs FP8 GEMM (General Matrix Multiplication) operation with optional native hardware support.
+    Args:
+        A (torch.Tensor): Input tensor A (FP8 or other dtype)
+        A_scale (torch.Tensor/float): Scale factor for tensor A
+        B (torch.Tensor): Input tensor B (FP8 or other dtype)
+        B_scale (torch.Tensor/float): Scale factor for tensor B
+        bias (torch.Tensor/None): Optional bias tensor
+        out_dtype (torch.dtype): Output data type
+        native_fp8_support (bool): Whether to use hardware-accelerated FP8 operations
+        
+    Returns:
+        torch.Tensor: Result of GEMM operation 
+    """
+    if A.numel() == 0:
+        # Deal with empty tensors (triggeted by empty MoE experts)
+        return torch.empty(size=(0, B.shape[0]), dtype=out_dtype, device=A.device)
+
+    if native_fp8_support:
+        need_reshape = A.dim() == 3
+        if need_reshape:
+            batch_size = A.shape[0]
+            A_input = A.reshape(-1, A.shape[-1])
+        else:
+            batch_size = None
+            A_input = A
+        output = torch._scaled_mm(
+            A_input,
+            B.t(),
+            out_dtype=out_dtype,
+            scale_a=A_scale,
+            scale_b=B_scale,
+            bias=bias,
+        )
+        if need_reshape:
+            output = output.reshape(
+                batch_size, output.shape[0] // batch_size, output.shape[1]
+            )
+    else:
+        output = torch.nn.functional.linear(
+            A.to(out_dtype) * A_scale,
+            B.to(out_dtype) * B_scale.to(out_dtype),
+            bias=bias,
+        )
+
+    return output
+
+def replace_module(model: torch.nn.Module, name: str, new_module: torch.nn.Module):
+    if "." in name:
+        parent_name = name.rsplit(".", 1)[0]
+        child_name = name[len(parent_name) + 1:]
+        parent = model.get_submodule(parent_name)
+    else:
+        parent_name = ""
+        parent = model
+        child_name = name
+    setattr(parent, child_name, new_module)
+
+
+# Class responsible for quantizing weights
+class FP8DynamicLinear(torch.nn.Module):
+    def __init__(
+            self,
+            weight: torch.Tensor,
+            weight_scale: torch.Tensor,
+            bias: torch.nn.Parameter,
+            native_fp8_support: bool = False,
+            dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__()
+        self.weight = torch.nn.Parameter(weight, requires_grad=False)
+        self.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+        self.bias = bias
+        self.native_fp8_support = native_fp8_support
+        self.dtype = dtype
+
+    # @torch.compile
+    def forward(self, x):
+        if x.dtype !=self.dtype:
+            x = x.to(self.dtype)
+        qinput, x_scale = per_tensor_quantize(x)
+        output = fp8_gemm(
+            A=qinput,
+            A_scale=x_scale,
+            B=self.weight,
+            B_scale=self.weight_scale,
+            bias=self.bias,
+            out_dtype=x.dtype,
+            native_fp8_support=self.native_fp8_support,
+        )
+        return output
+
+
+def FluxFp8GeMMProcessor(model: torch.nn.Module):
+    """Processes a PyTorch model to convert eligible Linear layers to FP8 precision.
+    
+    This function performs the following operations:
+    1. Checks for native FP8 support on the current GPU
+    2. Identifies target Linear layers in transformer blocks
+    3. Quantizes weights to FP8 format
+    4. Replaces original Linear layers with FP8DynamicLinear versions
+    5. Performs memory cleanup
+    
+    Args:
+        model (torch.nn.Module): The neural network model to be processed.
+                                Should contain transformer blocks with Linear layers.
+    """
+    native_fp8_support = (
+            torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0)
+    )
+    named_modules = list(model.named_modules())
+    for name, linear in tqdm.tqdm(named_modules, desc="Quantizing weights to fp8"):
+        if isinstance(linear, torch.nn.Linear) and "blocks" in name:
+            quant_weight, weight_scale = per_tensor_quantize(linear.weight)
+            bias = copy.deepcopy(linear.bias) if linear.bias is not None else None
+            quant_linear = FP8DynamicLinear(
+                weight=quant_weight,
+                weight_scale=weight_scale,
+                bias=bias,
+                native_fp8_support=native_fp8_support,
+                dtype=linear.weight.dtype
+            )
+            replace_module(model, name, quant_linear)
+            del linear.weight
+            del linear.bias
+            del linear
+    cleanup_memory()

requirements.txt

@@ -0,0 +1,19 @@
+torch==2.6.0
+torchvision==0.21.0
+triton==3.2.0
+transformers==4.57.0
+omegaconf==2.3.0
+ninja==1.13.0
+numpy==2.2.6
+einops==0.8.1
+moviepy==1.0.3
+opencv-python==4.12.0.88
+av==15.1.0
+plyfile==1.1.2
+ftfy==6.3.1
+flask==3.1.2
+gradio==5.49.1
+gsplat==1.5.2
+accelerate==1.10.1
+git+https://github.com/huggingface/diffusers.git@447e8322f76efea55d4769cd67c372edbf0715b8
+git+https://github.com/nerfstudio-project/gsplat.git@32f2a54d21c7ecb135320bb02b136b7407ae5712

utils.py

@@ -0,0 +1,531 @@
+from io import BytesIO
+import math
+import numpy as np
+import torch 
+import torch.nn as nn
+import torch.nn.functional as F
+import importlib
+from plyfile import PlyData, PlyElement
+
+import copy
+
+class EmbedContainer(nn.Module):
+    def __init__(self, tensor):
+        super().__init__()
+        self.tensor = nn.Parameter(tensor)
+    
+    def forward(self):
+        return self.tensor
+
+@torch.no_grad
+def zero_init(module):
+    if type(module) is torch.nn.Conv2d or type(module) is torch.nn.Linear:
+        module.weight.zero_()
+        module.bias.zero_()
+    return module
+
+def import_str(string):
+    # From https://github.com/CompVis/taming-transformers
+    module, cls = string.rsplit(".", 1)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+"""
+from https://github.com/Kai-46/minFM/blob/main/utils/ema.py
+Exponential Moving Average (EMA) utilities for PyTorch models.
+
+This module provides utilities for maintaining and updating EMA models,
+which are commonly used to improve model stability and generalization
+in training deep neural networks. It supports both regular tensors and
+DTensors (from FSDP-wrapped models).
+"""
+class EMA_FSDP:
+    def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999):
+        self.decay = decay
+        self.shadow = {}
+        self._init_shadow(fsdp_module)
+
+    @torch.no_grad()
+    def _init_shadow(self, fsdp_module):
+        # 判断是否是FSDP模型
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        if isinstance(fsdp_module, FSDP):
+            with FSDP.summon_full_params(fsdp_module, writeback=False):
+                for n, p in fsdp_module.module.named_parameters():
+                    self.shadow[n] = p.detach().clone().float().cpu()
+        else:
+            for n, p in fsdp_module.named_parameters():
+                self.shadow[n] = p.detach().clone().float().cpu()
+
+    @torch.no_grad()
+    def update(self, fsdp_module):
+        d = self.decay
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        if isinstance(fsdp_module, FSDP):
+            with FSDP.summon_full_params(fsdp_module, writeback=False):
+                for n, p in fsdp_module.module.named_parameters():
+                    self.shadow[n].mul_(d).add_(p.detach().float().cpu(), alpha=1. - d)
+        else:
+            for n, p in fsdp_module.named_parameters():
+                print(n, self.shadow[n])
+                self.shadow[n].mul_(d).add_(p.detach().float().cpu(), alpha=1. - d)
+
+    # Optional helpers ---------------------------------------------------
+    def state_dict(self):
+        return self.shadow            # picklable
+
+    def load_state_dict(self, sd):
+        self.shadow = {k: v.clone() for k, v in sd.items()}
+
+    def copy_to(self, fsdp_module):
+        # load EMA weights into an (unwrapped) copy of the generator
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=True):
+            for n, p in fsdp_module.module.named_parameters():
+                if n in self.shadow:
+                    p.data.copy_(self.shadow[n].to(p.dtype, device=p.device))
+
+def create_raymaps(cameras, h, w):
+    rays_o, rays_d = create_rays(cameras, h, w)
+    raymaps = torch.cat([rays_d, rays_o - (rays_o * rays_d).sum(dim=-1, keepdim=True) * rays_d], dim=-1)
+    return raymaps
+
+# def create_raymaps(cameras, h, w):
+#     rays_o, rays_d = create_rays(cameras, h, w)
+#     raymaps = torch.cat([rays_d, torch.cross(rays_d, rays_o, dim=-1)], dim=-1)
+#     return raymaps
+
+class EMANorm(nn.Module):
+    def __init__(self, beta):
+        super().__init__()
+        self.register_buffer('magnitude_ema', torch.ones([]))
+        self.beta = beta
+
+    def forward(self, x):
+        if self.training:
+            magnitude_cur = x.detach().to(torch.float32).square().mean()
+            self.magnitude_ema.copy_(magnitude_cur.lerp(self.magnitude_ema.to(torch.float32), self.beta))
+        input_gain = self.magnitude_ema.rsqrt()
+        x = x.mul(input_gain)
+        return x
+    
+class TimestepEmbedding(nn.Module):
+    def __init__(self, dim, max_period=10000, time_factor: float = 1000.0, zero_weight: bool = True):
+        super().__init__()
+        self.max_period = max_period
+        self.time_factor = time_factor
+        self.dim = dim
+        if zero_weight:
+            self.weight = nn.Parameter(torch.zeros(dim))
+        else:
+            self.weight = None
+
+    def forward(self, t):
+        if self.weight is None:
+            return timestep_embedding(t, self.dim, self.max_period, self.time_factor)
+        else:
+            return timestep_embedding(t, self.dim, self.max_period, self.time_factor) * self.weight.unsqueeze(0)
+
+@torch.compile(mode="max-autotune-no-cudagraphs", dynamic=True)
+def timestep_embedding(t, dim, max_period=10000, time_factor: float = 1000.0):
+    """
+    Create sinusoidal timestep embeddings.
+    :param t: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an (N, D) Tensor of positional embeddings.
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(t.device)
+
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+def quaternion_to_matrix(quaternions):
+    """
+    Convert rotations given as quaternions to rotation matrices.
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+# from https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html#matrix_to_quaternion
+def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+
+    Args:
+        quaternions: Quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Standardized quaternions as tensor of shape (..., 4).
+    """
+    return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
+
+def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
+    """
+    Returns torch.sqrt(torch.max(0, x))
+    but with a zero subgradient where x is 0.
+    """
+    ret = torch.zeros_like(x)
+    positive_mask = x > 0
+    if torch.is_grad_enabled():
+        ret[positive_mask] = torch.sqrt(x[positive_mask])
+    else:
+        ret = torch.where(positive_mask, torch.sqrt(x), ret)
+    return ret
+
+def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to quaternions.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+
+    batch_dim = matrix.shape[:-2]
+    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
+        matrix.reshape(batch_dim + (9,)), dim=-1
+    )
+
+    q_abs = _sqrt_positive_part(
+        torch.stack(
+            [
+                1.0 + m00 + m11 + m22,
+                1.0 + m00 - m11 - m22,
+                1.0 - m00 + m11 - m22,
+                1.0 - m00 - m11 + m22,
+            ],
+            dim=-1,
+        )
+    )
+
+    # we produce the desired quaternion multiplied by each of r, i, j, k
+    quat_by_rijk = torch.stack(
+        [
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
+        ],
+        dim=-2,
+    )
+
+    # We floor here at 0.1 but the exact level is not important; if q_abs is small,
+    # the candidate won't be picked.
+    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
+    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
+
+    # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
+    # forall i; we pick the best-conditioned one (with the largest denominator)
+    indices = q_abs.argmax(dim=-1, keepdim=True)
+    expand_dims = list(batch_dim) + [1, 4]
+    gather_indices = indices.unsqueeze(-1).expand(expand_dims)
+    out = torch.gather(quat_candidates, -2, gather_indices).squeeze(-2)
+    return standardize_quaternion(out)
+
+@torch.amp.autocast(device_type="cuda", enabled=False)
+def normalize_cameras(cameras, return_meta=False, ref_w2c=None, T_norm=None, n_frame=None):
+    B, N = cameras.shape[:2]
+     
+    c2ws = torch.zeros(B, N, 3, 4, device=cameras.device)
+
+    c2ws[..., :3, :3] = quaternion_to_matrix(cameras[..., 0:4])
+    c2ws[..., :, 3] = cameras[..., 4:7]
+
+    _c2ws = c2ws
+
+    ref_w2c = torch.inverse(matrix_to_square(_c2ws[:, :1])) if ref_w2c is None else ref_w2c
+    _c2ws = (ref_w2c.repeat(1, N, 1, 1) @ matrix_to_square(_c2ws))[..., :3, :]
+
+    if n_frame is not None:
+        T_norm = _c2ws[..., :n_frame, :3, 3].norm(dim=-1).max(dim=1)[0][..., None, None] if T_norm is None else T_norm
+    else:
+        T_norm = _c2ws[..., :3, 3].norm(dim=-1).max(dim=1)[0][..., None, None] if T_norm is None else T_norm
+
+    _c2ws[..., :3, 3] = _c2ws[..., :3, 3] / (T_norm + 1e-2)
+
+    R = matrix_to_quaternion(_c2ws[..., :3, :3])
+    T = _c2ws[..., :3, 3]
+    cameras = torch.cat([R.float(), T.float(), cameras[..., 7:]], dim=-1)
+
+    if return_meta:
+        return cameras, ref_w2c, T_norm
+    else:
+        return cameras
+
+def create_rays(cameras, h, w, uv_offset=None):
+    prefix_shape = cameras.shape[:-1]
+    cameras = cameras.flatten(0, -2)
+    device = cameras.device
+    N = cameras.shape[0]
+
+    c2w = torch.eye(4, device=device)[None].repeat(N, 1, 1)
+    c2w[:, :3, :3] = quaternion_to_matrix(cameras[:, :4])
+    c2w[:, :3, 3] = cameras[:, 4:7]
+
+    # fx, fy, cx, cy should be divided by original H, W
+    fx, fy, cx, cy = cameras[:, 7:].chunk(4, -1)
+
+    fx, cx = fx * w, cx * w
+    fy, cy = fy * h, cy * h
+
+    inds = torch.arange(0, h*w, device=device).expand(N, h*w)
+        
+    i = inds % w + 0.5
+    j = torch.div(inds, w, rounding_mode='floor') + 0.5
+
+    u = i / cx + (uv_offset[..., 0].reshape(N, h*w) if uv_offset is not None else 0) 
+    v = j / cy + (uv_offset[..., 1].reshape(N, h*w) if uv_offset is not None else 0) 
+
+    zs = - torch.ones_like(i)
+    xs = - (u - 1) * cx / fx * zs
+    ys = (v - 1) * cy / fy * zs
+    directions = torch.stack((xs, ys, zs), dim=-1)
+
+    rays_d = F.normalize(directions @ c2w[:, :3, :3].transpose(-1, -2), dim=-1)
+
+    rays_o = c2w[..., :3, 3] # [B, 3]
+    rays_o = rays_o[..., None, :].expand_as(rays_d)
+
+    rays_o = rays_o.reshape(*prefix_shape, h, w, 3)
+    rays_d = rays_d.reshape(*prefix_shape, h, w, 3)
+
+    return rays_o, rays_d
+
+def matrix_to_square(mat):
+    l = len(mat.shape)
+    if l==3:
+        return torch.cat([mat, torch.tensor([0,0,0,1]).repeat(mat.shape[0],1,1).to(mat.device)],dim=1)
+    elif l==4:
+        return torch.cat([mat, torch.tensor([0,0,0,1]).repeat(mat.shape[0],mat.shape[1],1,1).to(mat.device)],dim=2)
+
+def export_ply_for_gaussians(path, gaussians, opacity_threshold=0.00, T_norm=None):
+
+    sh_degree = int(math.sqrt((gaussians.shape[-1] - sum([3, 1, 3, 4])) / 3 - 1))
+
+    xyz, opacity, scale, rotation, feature = gaussians.float().split([3, 1, 3, 4, (sh_degree + 1)**2 * 3], dim=-1)
+     
+    means3D = xyz.contiguous().float()
+    opacity = opacity.contiguous().float()
+    scales = scale.contiguous().float()
+    rotations = rotation.contiguous().float()
+    shs = feature.contiguous().float() # [N, 1, 3]
+
+    # print(means3D.shape, opacity.shape, scales.shape, rotations.shape, shs.shape)
+
+    # prune by opacity
+    if opacity_threshold > 0:
+        mask = opacity[..., 0] >= opacity_threshold
+        means3D = means3D[mask]
+        opacity = opacity[mask]
+        scales = scales[mask]
+        rotations = rotations[mask]
+        shs = shs[mask]
+
+        print("Gaussian percentage: ", mask.float().mean())
+
+    if T_norm is not None:
+        means3D = means3D * T_norm.item()
+        scales = scales * T_norm.item()
+
+    # invert activation to make it compatible with the original ply format
+    opacity = torch.log(opacity/(1-opacity))
+    scales = torch.log(scales + 1e-8)
+
+    xyzs = means3D.detach() # .cpu().numpy()
+    f_dc = shs.detach().flatten(start_dim=1).contiguous() #.cpu().numpy()
+    opacities = opacity.detach() #.cpu().numpy()
+    scales = scales.detach() #.cpu().numpy()
+    rotations = rotations.detach() #.cpu().numpy()
+
+    l = ['x', 'y', 'z']
+    # All channels except the 3 DC
+    for i in range(f_dc.shape[1]):
+        l.append('f_dc_{}'.format(i))
+    l.append('opacity')
+    for i in range(scales.shape[1]):
+        l.append('scale_{}'.format(i))
+    for i in range(rotations.shape[1]):
+        l.append('rot_{}'.format(i))
+
+    dtype_full = [(attribute, 'f4') for attribute in l]
+
+    # 最优化方案：使用numpy的recarray直接创建
+    attributes = torch.cat((xyzs, f_dc, opacities, scales, rotations), dim=1).cpu().numpy()
+    
+    # 使用recarray直接创建，避免循环和类型转换
+    elements = np.rec.fromarrays([attributes[:, i] for i in range(attributes.shape[1])], names=l, formats=['f4'] * len(l))
+    el = PlyElement.describe(elements, 'vertex')
+
+    print(path)
+
+    PlyData([el]).write(path)
+
+    # plydata = PlyData([el])
+
+    # vert = plydata["vertex"]
+    # sorted_indices = np.argsort(
+    #     -np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"])
+    #     / (1 + np.exp(-vert["opacity"]))
+    # )
+    # buffer = BytesIO()
+    # for idx in sorted_indices:
+    #     v = plydata["vertex"][idx]
+    #     position = np.array([v["x"], v["y"], v["z"]], dtype=np.float32)
+    #     scales = np.exp(
+    #         np.array(
+    #             [v["scale_0"], v["scale_1"], v["scale_2"]],
+    #             dtype=np.float32,
+    #         )
+    #     )
+    #     rot = np.array(
+    #         [v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]],
+    #         dtype=np.float32,
+    #     )
+    #     SH_C0 = 0.28209479177387814
+    #     color = np.array(
+    #         [
+    #             0.5 + SH_C0 * v["f_dc_0"],
+    #             0.5 + SH_C0 * v["f_dc_1"],
+    #             0.5 + SH_C0 * v["f_dc_2"],
+    #             1 / (1 + np.exp(-v["opacity"])),
+    #         ]
+    #     )
+    #     buffer.write(position.tobytes())
+    #     buffer.write(scales.tobytes())
+    #     buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes())
+    #     buffer.write(
+    #         ((rot / np.linalg.norm(rot)) * 128 + 128)
+    #         .clip(0, 255)
+    #         .astype(np.uint8)
+    #         .tobytes()
+    #     )
+
+    # with open(path + '.splat', "wb") as f:
+    #     f.write(buffer.getvalue())
+
+@torch.amp.autocast(device_type="cuda", enabled=False)
+def quaternion_slerp(
+    q0, q1, fraction, spin: int = 0, shortestpath: bool = True
+):
+    """Return spherical linear interpolation between two quaternions.
+    Args:
+        quat0: first quaternion
+        quat1: second quaternion
+        fraction: how much to interpolate between quat0 vs quat1 (if 0, closer to quat0; if 1, closer to quat1)
+        spin: how much of an additional spin to place on the interpolation
+        shortestpath: whether to return the short or long path to rotation
+    """
+    d = (q0 * q1).sum(-1)
+    if shortestpath:
+        # invert rotation
+        d[d < 0.0] = -d[d < 0.0]
+        q1[d < 0.0] = q1[d < 0.0]
+
+    _d = d.clamp(0, 1.0)
+
+    # theta = torch.arccos(d) * fraction
+    # q2 = q1 - q0 * d
+    # q2 = q2 / (q2.norm(dim=-1) + 1e-10)
+    
+    # return torch.cos(theta) * q0 + torch.sin(theta) * q2
+
+    angle = torch.acos(_d) + spin * math.pi
+    isin = 1.0 / (torch.sin(angle)+ 1e-10)
+    q0_ = q0 * (torch.sin((1.0 - fraction) * angle) * isin)[..., None]
+    q1_ = q1 * (torch.sin(fraction * angle) * isin)[..., None]
+
+    q = q0_ + q1_
+
+    q[angle < 1e-5] = q0[angle < 1e-5]
+    # q[fraction < 1e-5] = q0[fraction < 1e-5]
+    # q[fraction > 1 - 1e-5] = q1[fraction > 1 - 1e-5]
+    # q[(d.abs() - 1).abs() < 1e-5] = q0[(d.abs() - 1).abs() < 1e-5]
+
+    return q
+
+def sample_from_two_pose(pose_a, pose_b, fraction, noise_strengths=[0, 0]):
+    """
+    Args:
+        pose_a: first pose
+        pose_b: second pose
+        fraction
+    """
+
+    quat_a = pose_a[..., :4]
+    quat_b = pose_b[..., :4]
+
+    dot = torch.sum(quat_a * quat_b, dim=-1, keepdim=True)
+    quat_b = torch.where(dot < 0, -quat_b, quat_b)
+
+    quaternion = quaternion_slerp(quat_a, quat_b, fraction)
+    quaternion = torch.nn.functional.normalize(quaternion + torch.randn_like(quaternion) * noise_strengths[0], dim=-1)
+
+    T = (1 - fraction)[:, None] * pose_a[..., 4:] + fraction[:, None] * pose_b[..., 4:]
+    T = T + torch.randn_like(T) * noise_strengths[1]
+
+    new_pose = pose_a.clone()
+    new_pose[..., :4] = quaternion
+    new_pose[..., 4:] = T
+    return new_pose
+
+def sample_from_dense_cameras(dense_cameras, t, noise_strengths=[0, 0, 0, 0]):
+    N, C = dense_cameras.shape
+    M = t.shape
+    
+    left = torch.floor(t * (N-1)).long().clamp(0, N-2)
+    right = left + 1
+    fraction = t * (N-1) - left
+
+    a = torch.gather(dense_cameras, 0, left[..., None].repeat(1, C))
+    b = torch.gather(dense_cameras, 0, right[..., None].repeat(1, C))
+
+    new_pose = sample_from_two_pose(a[:, :7], 
+                                    b[:, :7], fraction, noise_strengths=noise_strengths[:2])
+
+    new_ins = (1 - fraction)[:, None] * a[:, 7:] + fraction[:, None] * b[:, 7:]
+
+    return torch.cat([new_pose, new_ins], dim=1)
+