File size: 3,569 Bytes
d6b621b 4d01bea d6b621b 4d01bea d6b621b 4d01bea d6b621b 4d01bea d6b621b 4d01bea d6b621b 4d01bea a9d5db4 d6b621b 4d01bea d6b621b 4d01bea 20c712c d6b621b c9fe23b 2070c2d d44fe94 a0e034f c9fe23b d44fe94 d6b621b c9fe23b d6b621b 20c712c | 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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | from typing import Dict, List, Any
import json
import numpy as np
from transformers import AutoProcessor, MusicgenForConditionalGeneration
import torch
import logging
class EndpointHandler:
def __init__(self, path=""):
# load model and processor from path
self.processor = AutoProcessor.from_pretrained(path)
# Check if CUDA is available, and set the device accordingly
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load the model to the device
self.model = MusicgenForConditionalGeneration.from_pretrained(path)
self.model.to(self.device) # Correcting this line
def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
"""
Args:
data (:dict:):
The payload with the text prompt and generation parameters.
"""
# Set up logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
# process input
logger.debug(f"Data: {data}")
inputs = data.pop("inputs", data)
logger.debug(f"Inputs: {inputs}")
parameters = data.pop("parameters", None)
logger.debug(f"Parameters: {parameters}")
duration = parameters.pop("duration", None)
logger.debug(f"Duration: {duration}")
is_existing_audio = parameters.pop("is_existing_audio", None)
audio = parameters.pop("audio", None)
logger.debug(f"Audio: {audio}")
sampling_rate = parameters.pop("sampling_rate", None)
logger.debug(f"Sampling Rate: {sampling_rate}")
if not sampling_rate:
sampling_rate = self.model.config.audio_encoder.sampling_rate
if audio is not None and is_existing_audio == False:
audio_array = np.array(audio)
audio = audio_array[: len(audio_array) // 3]
# sample["array"] = sample["array"][: len(sample["array"]) // 2]
elif audio is not None and is_existing_audio == True:
audio_array = np.array(audio)
audio = audio_array[len(audio_array) // 3:]
if duration is not None:
# Calculate max new tokens based on duration, this is a placeholder, replace with actual logic
max_new_tokens = int(duration * 50)
else:
max_new_tokens = 256 # Default value if duration is not provided
# preprocess
inputs = self.processor(
text=[inputs],
padding=True,
return_tensors="pt",
audio=audio,
sampling_rate=sampling_rate).to(self.device)
# If 'duration' is inside 'parameters', remove it
if parameters is not None and 'duration' in parameters:
parameters.pop('duration')
if parameters is not None and 'audio' in parameters:
parameters.pop('audio')
if parameters is not None and 'sampling_rate' in parameters:
parameters.pop('sampling_rate')
if parameters is not None and 'is_existing_audio' in parameters:
parameters.pop('is_existing_audio')
# pass inputs with all kwargs in data
if parameters is not None:
outputs = self.model.generate(**inputs, max_new_tokens=max_new_tokens, **parameters)
else:
outputs = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
# postprocess the prediction
prediction = outputs[0].cpu().numpy()
return [{"generated_text": prediction, "sampling_rate" : sampling_rate}] |