| |
| """ |
| 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") |
|
|
