init

app.py

@@ -28,6 +28,15 @@ local_path = snapshot_download(
     revision="main",                  
 )
 
+local_path_vae = snapshot_download(
+    repo_id="wusize/Puffin",
+    repo_type="model",       
+    #filename="Puffin-Base.pth",         
+    local_dir="checkpoints_vae/",  
+    local_dir_use_symlinks=False,     
+    revision="main",                  
+)
+
 
 NUM = r"[+-]?(?:\d+(?:\.\d+)?|\.\d+)(?:[eE][+-]?\d+)?"
 CAM_PATTERN = re.compile(r"(?:camera parameters.*?:|roll.*?:)\s*("+NUM+r")\s*,\s*("+NUM+r")\s*,\s*("+NUM+r")", re.IGNORECASE|re.DOTALL)
@@ -48,6 +57,10 @@ checkpoint_path = "checkpoints/Puffin-Base.pth"
 checkpoint = torch.load(checkpoint_path)
 info = model.load_state_dict(checkpoint, strict=False)
 
+checkpoint_path_vae = "checkpoints_vae/vae.pth"
+checkpoint_vae = torch.load(checkpoint_path_vae)
+info_vae = model.vae.load_state_dict(checkpoint_vae, strict=False)
+
 def fig_to_image(fig):
     buf = io.BytesIO()
     fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)