init

app.py

@@ -7,6 +7,7 @@ import math
 import re
 from einops import rearrange
 from mmengine.config import Config
+from xtuner.registry import BUILDER 
 
 import matplotlib
 matplotlib.use("Agg")
@@ -15,10 +16,6 @@ import matplotlib.pyplot as plt
 from scripts.camera.cam_dataset import Cam_Generator
 from scripts.camera.visualization.visualize_batch import make_perspective_figures
 
-from mmengine.registry import Registry
-__all__ = ['BUILDER']
-BUILDER = Registry('builder')
-
 from huggingface_hub import snapshot_download
 import os
 local_path = snapshot_download(
@@ -45,16 +42,11 @@ def center_crop(image):
 ##### load model
 config = "configs/pipelines/stage_2_base.py"
 config = Config.fromfile(config)
-model = BUILDER.build(config.model).eval()
+model = BUILDER.build(config.model).cuda().bfloat16().eval()
 checkpoint_path = "checkpoints/Puffin-Base.pth"
 checkpoint = torch.load(checkpoint_path)
 info = model.load_state_dict(checkpoint, strict=False)
 
-if torch.cuda.is_available():
-    model = model.to(torch.bfloat16).cuda()
-else:
-    model = model.to(torch.float32)
-
 
 @torch.inference_mode()
 @spaces.GPU(duration=120)