new

app.py

@@ -482,7 +482,7 @@ def main():
 
     # demo.queue(default_concurrency_limit=1)
     # demo.launch(server_name="0.0.0.0", server_port=9239)
-    subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+    # subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
     demo.launch()
 
 

models/llm/llm.py

@@ -1,6 +1,7 @@
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import spaces
 # device = "cuda" # the device to load the model onto
 model_name_or_dir="Qwen/Qwen2-7B-Instruct"
 
@@ -55,7 +56,7 @@ def load_llm_model(model_name_or_dir, torch_dtype='auto', device_map='cpu'):
 # print(f"Before load llm model: {torch.cuda.memory_allocated() / 1024**3} GB")
 # load_model()
 # print(f"After load llm model: {torch.cuda.memory_allocated() / 1024**3} GB")
-
+@spaces.GPU
 def get_llm_response(model, tokenizer, user_prompt, seed=None, system_prompt=DEFAULT_SYSTEM_PROMPT):
     # global model
     # global tokenizer

pipeline/kiss3d_wrapper.py

@@ -256,7 +256,7 @@ class kiss3d_wrapper(object):
         caption_text = self.get_detailed_prompt(caption_text)
 
         return caption_text
-
+    @spaces.GPU
     def get_detailed_prompt(self, prompt, seed=None):
         if self.llm_model is not None:
             detailed_prompt = get_llm_response(self.llm_model, self.llm_tokenizer, prompt, seed=seed)