File size: 1,489 Bytes
1800268 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
#!/usr/bin/env python
"""
Load and use the 8-bit quantized VibeVoice model
"""
import torch
from transformers import BitsAndBytesConfig
from vibevoice.modular.modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference
from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor
def load_quantized_model(model_path="/kaggle/working/quantized_8bit_FIXED"):
"""Load the pre-quantized VibeVoice model"""
print("Loading 8-bit quantized VibeVoice model...")
bnb_config = BitsAndBytesConfig(
load_in_8bit=True,
bnb_8bit_compute_dtype=torch.bfloat16,
)
processor = VibeVoiceProcessor.from_pretrained(model_path)
model = VibeVoiceForConditionalGenerationInference.from_pretrained(
model_path,
quantization_config=bnb_config,
device_map='cuda',
torch_dtype=torch.bfloat16,
)
model.eval()
print("✅ Model loaded successfully!")
print(f"💾 Memory usage: {torch.cuda.memory_allocated() / 1e9:.1f} GB")
return model, processor
if __name__ == "__main__":
model, processor = load_quantized_model()
text = "Speaker 1: Hello! Speaker 2: Hi there!"
inputs = processor(
text=[text],
voice_samples=[["path/to/voice1.wav", "path/to/voice2.wav"]],
padding=True,
return_tensors="pt",
)
with torch.no_grad():
outputs = model.generate(**inputs)
processor.save_audio(outputs.speech_outputs[0], "output.wav")
|