updating for zerogpu

app.py

@@ -19,6 +19,7 @@ from pathlib import Path
 import uuid
 import shutil
 from huggingface_hub import snapshot_download
+import spaces
 
 
 GLOBAL_STATE = {}
@@ -91,6 +92,25 @@ def build_interface():
                 
     return interface
 
+@spaces.GPU
+def generate(args):
+    config_path = args.config_path
+    gpt_ckpt_path = "./model_weights/shape_gpt.safetensors"
+    shape_ckpt_path = "./model_weights/shape_tokenizer.safetensors"
+    engine_fast = EngineFast(
+        config_path,
+        gpt_ckpt_path, 
+        shape_ckpt_path,
+        device=torch.device("cuda"),
+    )
+    GLOBAL_STATE["engine_fast"] = engine_fast
+    GLOBAL_STATE["SAVE_DIR"] = args.save_dir
+    os.makedirs(GLOBAL_STATE["SAVE_DIR"], exist_ok=True)
+
+    demo = build_interface()
+    demo.queue(default_concurrency_limit=1)
+    demo.launch()
+
 if __name__=="__main__":
 
     parser = argparse.ArgumentParser()
@@ -123,19 +143,4 @@ if __name__=="__main__":
         repo_id="Roblox/cube3d-v0.1",
         local_dir="./model_weights"
     )
-    config_path = args.config_path
-    gpt_ckpt_path = "./model_weights/shape_gpt.safetensors"
-    shape_ckpt_path = "./model_weights/shape_tokenizer.safetensors"
-    engine_fast = EngineFast(
-        config_path,
-        gpt_ckpt_path, 
-        shape_ckpt_path,
-        device=torch.device("cuda"),
-    )
-    GLOBAL_STATE["engine_fast"] = engine_fast
-    GLOBAL_STATE["SAVE_DIR"] = args.save_dir
-    os.makedirs(GLOBAL_STATE["SAVE_DIR"], exist_ok=True)
-
-    demo = build_interface()
-    demo.queue(default_concurrency_limit=1)
-    demo.launch()
+    generate(args)