rajkumarrawal commited on
Commit
826cc00
·
1 Parent(s): 6f9b46a

refactor: simplify model device handling using accelerate library

Browse files
Files changed (2) hide show
  1. app.py +3 -3
  2. requirements.txt +1 -0
app.py CHANGED
@@ -12,11 +12,11 @@ model_name = 'Marqo/marqo-fashionSigLIP'
12
 
13
  # Force CPU usage to avoid device mapping issues
14
  device = torch.device('cpu')
 
 
15
  model = AutoModel.from_pretrained(
16
  model_name,
17
- trust_remote_code=True,
18
- torch_dtype=torch.float32,
19
- device_map=None # This forces the model to stay on CPU
20
  ).to(device)
21
 
22
  processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
 
12
 
13
  # Force CPU usage to avoid device mapping issues
14
  device = torch.device('cpu')
15
+
16
+ # Simple loading approach - let model handle device placement
17
  model = AutoModel.from_pretrained(
18
  model_name,
19
+ trust_remote_code=True
 
 
20
  ).to(device)
21
 
22
  processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
requirements.txt CHANGED
@@ -4,6 +4,7 @@ requests
4
  Pillow
5
  open_clip_torch
6
  ftfy
 
7
 
8
  # This is only needed for local deployment
9
  gradio
 
4
  Pillow
5
  open_clip_torch
6
  ftfy
7
+ accelerate
8
 
9
  # This is only needed for local deployment
10
  gradio