bachvudinh commited on
Commit
39cc970
1 Parent(s): e10af0d

add load model outside a function

Browse files
Files changed (1) hide show
  1. app.py +21 -32
app.py CHANGED
@@ -20,6 +20,7 @@ vq_model = RQBottleneckTransformer.load_model(
20
  "whisper-vq-stoks-medium-en+pl-fixed.model"
21
  ).to(device)
22
  vq_model.ensure_whisper(device)
 
23
  @spaces.GPU
24
  def audio_to_sound_tokens_whisperspeech(audio_path):
25
  wav, sr = torchaudio.load(audio_path)
@@ -31,6 +32,7 @@ def audio_to_sound_tokens_whisperspeech(audio_path):
31
 
32
  result = ''.join(f'<|sound_{num:04d}|>' for num in codes)
33
  return f'<|sound_start|>{result}<|sound_end|>'
 
34
  @spaces.GPU
35
  def audio_to_sound_tokens_whisperspeech_transcribe(audio_path):
36
  wav, sr = torchaudio.load(audio_path)
@@ -42,45 +44,28 @@ def audio_to_sound_tokens_whisperspeech_transcribe(audio_path):
42
 
43
  result = ''.join(f'<|sound_{num:04d}|>' for num in codes)
44
  return f'<|reserved_special_token_69|><|sound_start|>{result}<|sound_end|>'
45
- def audio_to_sound_tokens(audio_path, target_bandwidth=1.5, device="cuda"):
46
- model = EncodecModel.encodec_model_24khz()
47
- model.set_target_bandwidth(target_bandwidth)
48
- model.to(device)
49
-
50
- wav, sr = torchaudio.load(audio_path)
51
- wav = convert_audio(wav, sr, model.sample_rate, model.channels)
52
- wav = wav.unsqueeze(0).to(device)
53
-
54
- with torch.no_grad():
55
- encoded_frames = model.encode(wav)
56
- codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1)
57
-
58
- audio_code1, audio_code2 = codes[0][0], codes[0][1]
59
- flatten_tokens = torch.stack((audio_code1, audio_code2), dim=1).flatten().tolist()
60
- result = ''.join(f'<|sound_{num:04d}|>' for num in flatten_tokens)
61
- return f'<|sound_start|>{result}<|sound_end|>'
62
- @spaces.GPU
63
- def setup_pipeline(model_path, use_4bit=False, use_8bit=False):
64
- tokenizer = AutoTokenizer.from_pretrained(model_path)
65
- model_kwargs = {"device_map": "auto"}
66
- if use_8bit:
67
- model_kwargs["quantization_config"] = BitsAndBytesConfig(
68
- load_in_8bit=True,
69
- llm_int8_enable_fp32_cpu_offload=False,
70
- llm_int8_has_fp16_weight=False,
71
- )
72
- else:
73
- model_kwargs["torch_dtype"] = torch.bfloat16
74
- model = AutoModelForCausalLM.from_pretrained(model_path, **model_kwargs)
75
- return pipeline("text-generation", model=model, tokenizer=tokenizer)
76
 
77
  tts = TTSProcessor(device)
 
78
  llm_path = "homebrewltd/Llama3.1-s-instruct-2024-08-19-epoch-3"
79
- pipe = setup_pipeline(llm_path, use_8bit=False)
 
 
 
 
 
 
 
 
 
 
 
80
  tokenizer = pipe.tokenizer
81
  model = pipe.model
82
  # print(tokenizer.encode("<|sound_0001|>", add_special_tokens=False))# return the audio tensor
83
  # print(tokenizer.eos_token)
 
 
84
  @spaces.GPU
85
  def text_to_audio_file(text):
86
  # gen a random id for the audio file
@@ -96,6 +81,8 @@ def text_to_audio_file(text):
96
  # torchaudio.save(temp_file, audio.cpu(), sample_rate=24000)
97
  print(f"Saved audio to {temp_file}")
98
  return temp_file
 
 
99
  @spaces.GPU
100
  def process_input(input_type, text_input=None, audio_file=None):
101
  # if input_type == "text":
@@ -106,6 +93,7 @@ def process_input(input_type, text_input=None, audio_file=None):
106
 
107
  # if input_type == "text":
108
  # os.remove(audio_file)
 
109
  @spaces.GPU
110
  def process_transcribe_input(input_type, text_input=None, audio_file=None):
111
  # if input_type == "text":
@@ -124,6 +112,7 @@ class StopOnTokens(StoppingCriteria):
124
  if input_ids[0][-1] == stop_id:
125
  return True
126
  return False
 
127
  @spaces.GPU
128
  def process_audio(audio_file, transcript=False):
129
  if audio_file is None:
 
20
  "whisper-vq-stoks-medium-en+pl-fixed.model"
21
  ).to(device)
22
  vq_model.ensure_whisper(device)
23
+
24
  @spaces.GPU
25
  def audio_to_sound_tokens_whisperspeech(audio_path):
26
  wav, sr = torchaudio.load(audio_path)
 
32
 
33
  result = ''.join(f'<|sound_{num:04d}|>' for num in codes)
34
  return f'<|sound_start|>{result}<|sound_end|>'
35
+
36
  @spaces.GPU
37
  def audio_to_sound_tokens_whisperspeech_transcribe(audio_path):
38
  wav, sr = torchaudio.load(audio_path)
 
44
 
45
  result = ''.join(f'<|sound_{num:04d}|>' for num in codes)
46
  return f'<|reserved_special_token_69|><|sound_start|>{result}<|sound_end|>'
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
47
 
48
  tts = TTSProcessor(device)
49
+ use_8bit = False
50
  llm_path = "homebrewltd/Llama3.1-s-instruct-2024-08-19-epoch-3"
51
+ tokenizer = AutoTokenizer.from_pretrained(llm_path)
52
+ model_kwargs = {"device_map": "auto"}
53
+ if use_8bit:
54
+ model_kwargs["quantization_config"] = BitsAndBytesConfig(
55
+ load_in_8bit=True,
56
+ llm_int8_enable_fp32_cpu_offload=False,
57
+ llm_int8_has_fp16_weight=False,
58
+ )
59
+ else:
60
+ model_kwargs["torch_dtype"] = torch.bfloat16
61
+ model = AutoModelForCausalLM.from_pretrained(llm_path, **model_kwargs)
62
+ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
63
  tokenizer = pipe.tokenizer
64
  model = pipe.model
65
  # print(tokenizer.encode("<|sound_0001|>", add_special_tokens=False))# return the audio tensor
66
  # print(tokenizer.eos_token)
67
+
68
+
69
  @spaces.GPU
70
  def text_to_audio_file(text):
71
  # gen a random id for the audio file
 
81
  # torchaudio.save(temp_file, audio.cpu(), sample_rate=24000)
82
  print(f"Saved audio to {temp_file}")
83
  return temp_file
84
+
85
+
86
  @spaces.GPU
87
  def process_input(input_type, text_input=None, audio_file=None):
88
  # if input_type == "text":
 
93
 
94
  # if input_type == "text":
95
  # os.remove(audio_file)
96
+
97
  @spaces.GPU
98
  def process_transcribe_input(input_type, text_input=None, audio_file=None):
99
  # if input_type == "text":
 
112
  if input_ids[0][-1] == stop_id:
113
  return True
114
  return False
115
+
116
  @spaces.GPU
117
  def process_audio(audio_file, transcript=False):
118
  if audio_file is None: