handling zerogpu in prompt handling

app.py

@@ -1,3 +1,4 @@
+import spaces
 import argparse
 import gradio as gr
 import os
@@ -19,7 +20,6 @@ from pathlib import Path
 import uuid
 import shutil
 from huggingface_hub import snapshot_download
-import spaces
 
 
 GLOBAL_STATE = {}
@@ -40,8 +40,22 @@ def gen_save_folder(max_size=200):
 
     return new_folder
 
+@spaces.GPU
 def handle_text_prompt(input_prompt, variance = 0):
     print(f"prompt: {input_prompt}, variance: {variance}")
+
+    if "engine_fast" not in GLOBAL_STATE: 
+        config_path = GLOBAL_STATE["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
+
     top_p = None if variance == 0 else (100 - variance) / 100.0
     mesh_v_f = GLOBAL_STATE["engine_fast"].t2s([input_prompt], use_kv_cache=True, resolution_base=8.0, top_p=top_p)
     # save output
@@ -92,18 +106,8 @@ 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["config_path"] = args.config_path
     GLOBAL_STATE["SAVE_DIR"] = args.save_dir
     os.makedirs(GLOBAL_STATE["SAVE_DIR"], exist_ok=True)