merve HF Staff commited on
Commit
78fcdc5
·
verified ·
1 Parent(s): a1b7c96

carry model within decorator

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -7,7 +7,6 @@ from transformers import AutoImageProcessor, AutoModel
7
  import torch.nn.functional as F
8
  import spaces
9
 
10
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
 
12
  DINO_MODELS = {
13
  "DINOv3 Base ViT": "facebook/dinov3-vitb16-pretrain-lvd1689m",
@@ -22,13 +21,15 @@ def load_model(model_name):
22
 
23
  processor = AutoImageProcessor.from_pretrained(model_path)
24
  model = AutoModel.from_pretrained(model_path)
25
- model = model.to(device)
26
  return f"✅ Model '{model_name}' loaded successfully!"
27
 
28
  load_model("DINOv3 Base ViT")
29
 
30
  @spaces.GPU()
31
  def extract_features(image):
 
 
 
32
  original_size = image.size
33
  inputs = processor(images=image, return_tensors="pt")
34
  inputs = {k: v.to(device) for k, v in inputs.items()}
 
7
  import torch.nn.functional as F
8
  import spaces
9
 
 
10
 
11
  DINO_MODELS = {
12
  "DINOv3 Base ViT": "facebook/dinov3-vitb16-pretrain-lvd1689m",
 
21
 
22
  processor = AutoImageProcessor.from_pretrained(model_path)
23
  model = AutoModel.from_pretrained(model_path)
 
24
  return f"✅ Model '{model_name}' loaded successfully!"
25
 
26
  load_model("DINOv3 Base ViT")
27
 
28
  @spaces.GPU()
29
  def extract_features(image):
30
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
31
+
32
+ model = model.to(device)
33
  original_size = image.size
34
  inputs = processor(images=image, return_tensors="pt")
35
  inputs = {k: v.to(device) for k, v in inputs.items()}