cbensimon HF Staff commited on
Commit
4f27510
·
1 Parent(s): 060b6b6

Worker re-use

Browse files
Files changed (1) hide show
  1. app.py +6 -2
app.py CHANGED
@@ -57,10 +57,14 @@ weights = {name: tensor.to('cuda') for name, tensor in weights.items()}
57
 
58
  del model
59
 
 
 
60
  @spaces.GPU
61
  def run_model():
62
- compiled_model: AOTICompiledModel = torch._inductor.aoti_load_package(package_path)
63
- compiled_model.load_constants(weights, check_full_update=True, user_managed=True)
 
 
64
  with torch.inference_mode():
65
  return str(compiled_model(example_inputs))
66
 
 
57
 
58
  del model
59
 
60
+ compiled_model: AOTICompiledModel | None = None
61
+
62
  @spaces.GPU
63
  def run_model():
64
+ global compiled_model
65
+ if compiled_model is None:
66
+ compiled_model = torch._inductor.aoti_load_package(package_path)
67
+ compiled_model.load_constants(weights, check_full_update=True, user_managed=True)
68
  with torch.inference_mode():
69
  return str(compiled_model(example_inputs))
70