MohamedRashad commited on
Commit
aa14541
·
1 Parent(s): 6487ffb

Refactor model loading to use direct model paths instead of snapshot downloads

Browse files
Files changed (1) hide show
  1. app.py +6 -26
app.py CHANGED
@@ -1,37 +1,17 @@
1
- from pathlib import Path
2
-
3
  import gradio as gr
4
  import spaces
5
  import torch
6
- from huggingface_hub import snapshot_download
7
  from transformers import AutoProcessor, VoxtralForConditionalGeneration
8
 
9
- # Model paths and setup
10
- voxtral_mini_path = snapshot_download(
11
- repo_id='mistralai/Voxtral-Mini-3B-2507',
12
- revision='refs/pr/16',
13
- local_dir=Path(__file__).parent / 'Voxtral-Mini-3B-2507',
14
- resume_download=True,
15
- )
16
- print(f"Voxtral Mini model downloaded to: {voxtral_mini_path}")
17
-
18
- voxtral_small_path = snapshot_download(
19
- repo_id='mistralai/Voxtral-Small-24B-2507',
20
- revision='refs/pr/9',
21
- local_dir=Path(__file__).parent / 'Voxtral-Small-24B-2507',
22
- resume_download=True,
23
- )
24
- print(f"Voxtral Small model downloaded to: {voxtral_small_path}")
25
-
26
  device = "cuda" if torch.cuda.is_available() else "cpu"
27
  print(f"Using device: {device}")
28
 
29
  # Load model and processor
30
- voxtral_mini_processor = AutoProcessor.from_pretrained(voxtral_mini_path)
31
- voxtral_mini_model = VoxtralForConditionalGeneration.from_pretrained(voxtral_mini_path, torch_dtype=torch.bfloat16, device_map=device)
32
 
33
- voxtral_small_processor = AutoProcessor.from_pretrained(voxtral_small_path)
34
- voxtral_small_model = VoxtralForConditionalGeneration.from_pretrained(voxtral_small_path, torch_dtype=torch.bfloat16, device_map=device)
35
 
36
  @spaces.GPU()
37
  def process_audio(audio_path, model_name, language="en", max_tokens=500):
@@ -42,11 +22,11 @@ def process_audio(audio_path, model_name, language="en", max_tokens=500):
42
  if model_name == "Voxtral Mini (3B)":
43
  model = voxtral_mini_model
44
  processor = voxtral_mini_processor
45
- repo_id = str(voxtral_mini_path)
46
  elif model_name == "Voxtral Small (24B)":
47
  model = voxtral_small_model
48
  processor = voxtral_small_processor
49
- repo_id = str(voxtral_small_path)
50
  else:
51
  return "Invalid model selected."
52
 
 
 
 
1
  import gradio as gr
2
  import spaces
3
  import torch
 
4
  from transformers import AutoProcessor, VoxtralForConditionalGeneration
5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
  device = "cuda" if torch.cuda.is_available() else "cpu"
7
  print(f"Using device: {device}")
8
 
9
  # Load model and processor
10
+ voxtral_mini_processor = AutoProcessor.from_pretrained("MohamedRashad/Voxtral-Mini-3B-2507-transformers")
11
+ voxtral_mini_model = VoxtralForConditionalGeneration.from_pretrained("MohamedRashad/Voxtral-Mini-3B-2507-transformers", torch_dtype=torch.bfloat16, device_map=device)
12
 
13
+ voxtral_small_processor = AutoProcessor.from_pretrained("MohamedRashad/Voxtral-Small-24B-2507-transformers")
14
+ voxtral_small_model = VoxtralForConditionalGeneration.from_pretrained("MohamedRashad/Voxtral-Small-24B-2507-transformers", torch_dtype=torch.bfloat16, device_map=device)
15
 
16
  @spaces.GPU()
17
  def process_audio(audio_path, model_name, language="en", max_tokens=500):
 
22
  if model_name == "Voxtral Mini (3B)":
23
  model = voxtral_mini_model
24
  processor = voxtral_mini_processor
25
+ repo_id = "MohamedRashad/Voxtral-Mini-3B-2507-transformers"
26
  elif model_name == "Voxtral Small (24B)":
27
  model = voxtral_small_model
28
  processor = voxtral_small_processor
29
+ repo_id = "MohamedRashad/Voxtral-Small-24B-2507-transformers"
30
  else:
31
  return "Invalid model selected."
32