p1atdev commited on
Commit
23c73d4
1 Parent(s): 4d398f3

fix: model device

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -16,8 +16,8 @@ MODEL_NAME = MODEL_NAMES[0]
16
 
17
  model = AutoModelForImageClassification.from_pretrained(
18
  MODEL_NAME,
19
- device="cuda",
20
  )
 
21
  processor = AutoImageProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
22
 
23
 
@@ -90,7 +90,7 @@ def animagine_prompt(rating: list[str], character: list[str], general: list[str]
90
  return ", ".join(all_tags)
91
 
92
 
93
- @spaces.GPU
94
  def predict_tags(
95
  image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8
96
  ):
 
16
 
17
  model = AutoModelForImageClassification.from_pretrained(
18
  MODEL_NAME,
 
19
  )
20
+ model.to("cuda" if torch.cuda.is_available() else "cpu")
21
  processor = AutoImageProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
22
 
23
 
 
90
  return ", ".join(all_tags)
91
 
92
 
93
+ @spaces.GPU(enable_queue=True)
94
  def predict_tags(
95
  image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8
96
  ):