shimu0215 commited on
Commit
facbd60
·
verified ·
1 Parent(s): f4b8ef8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -4
app.py CHANGED
@@ -9,6 +9,7 @@ from torchvision import transforms
9
  from PIL import Image
10
  import matplotlib.pyplot as plt
11
  import tempfile
 
12
 
13
  class Config:
14
  ASSETS_DIR = os.path.join(os.path.dirname(__file__), 'assets')
@@ -28,10 +29,11 @@ class Config:
28
  class ModelManager:
29
  @staticmethod
30
  def load_model(checkpoint_name: str):
31
- if checkpoint_name is None:
32
- return None
33
- checkpoint_path = os.path.join(Config.CHECKPOINTS_DIR, checkpoint_name)
34
- model = torch.jit.load(checkpoint_path)
 
35
  model.eval()
36
  model.to("cuda")
37
  return model
 
9
  from PIL import Image
10
  import matplotlib.pyplot as plt
11
  import tempfile
12
+ from huggingface_hub import hf_hub_download
13
 
14
  class Config:
15
  ASSETS_DIR = os.path.join(os.path.dirname(__file__), 'assets')
 
29
  class ModelManager:
30
  @staticmethod
31
  def load_model(checkpoint_name: str):
32
+ model_path = hf_hub_download(
33
+ repo_id="shimu0215/seg",
34
+ filename="sapiens_2b_render_people_epoch_25_torchscript.pt2",
35
+ )
36
+ model = torch.jit.load(model_path)
37
  model.eval()
38
  model.to("cuda")
39
  return model