HuggingGPT / models_server.py
Shenziqian's picture
Duplicate from microsoft/HuggingGPT
a5fb637
import argparse
import logging
import random
import uuid
import numpy as np
from transformers import pipeline
from diffusers import DiffusionPipeline, StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
from diffusers.utils import load_image
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.utils import export_to_video
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5ForSpeechToSpeech
from transformers import BlipProcessor, BlipForConditionalGeneration
from transformers import TrOCRProcessor, VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from datasets import load_dataset
from PIL import Image
import io
from torchvision import transforms
import torch
import torchaudio
from speechbrain.pretrained import WaveformEnhancement
import joblib
from huggingface_hub import hf_hub_url, cached_download
from transformers import AutoImageProcessor, TimesformerForVideoClassification
from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation, AutoFeatureExtractor
from controlnet_aux import OpenposeDetector, MLSDdetector, HEDdetector, CannyDetector, MidasDetector
from controlnet_aux.open_pose.body import Body
from controlnet_aux.mlsd.models.mbv2_mlsd_large import MobileV2_MLSD_Large
from controlnet_aux.hed import Network
from transformers import DPTForDepthEstimation, DPTFeatureExtractor
import warnings
import time
from espnet2.bin.tts_inference import Text2Speech
import soundfile as sf
from asteroid.models import BaseModel
import traceback
import os
import yaml
warnings.filterwarnings("ignore")
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="config.yaml")
args = parser.parse_args()
if __name__ != "__main__":
args.config = "config.gradio.yaml"
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
config = yaml.load(open(args.config, "r"), Loader=yaml.FullLoader)
local_deployment = config["local_deployment"]
if config["inference_mode"] == "huggingface":
local_deployment = "none"
PROXY = None
if config["proxy"]:
PROXY = {
"https": config["proxy"],
}
start = time.time()
# local_models = "models/"
local_models = ""
def load_pipes(local_deployment):
other_pipes = {}
standard_pipes = {}
controlnet_sd_pipes = {}
if local_deployment in ["full"]:
other_pipes = {
# "Salesforce/blip-image-captioning-large": {
# "model": BlipForConditionalGeneration.from_pretrained(f"Salesforce/blip-image-captioning-large"),
# "processor": BlipProcessor.from_pretrained(f"Salesforce/blip-image-captioning-large"),
# "device": "cuda:0"
# },
"damo-vilab/text-to-video-ms-1.7b": {
"model": DiffusionPipeline.from_pretrained(f"{local_models}damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16, variant="fp16"),
"device": "cuda:0"
},
# "facebook/maskformer-swin-large-ade": {
# "model": MaskFormerForInstanceSegmentation.from_pretrained(f"facebook/maskformer-swin-large-ade"),
# "feature_extractor" : AutoFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-ade"),
# "device": "cuda:0"
# },
# "microsoft/trocr-base-printed": {
# "processor": TrOCRProcessor.from_pretrained(f"microsoft/trocr-base-printed"),
# "model": VisionEncoderDecoderModel.from_pretrained(f"microsoft/trocr-base-printed"),
# "device": "cuda:0"
# },
# "microsoft/trocr-base-handwritten": {
# "processor": TrOCRProcessor.from_pretrained(f"microsoft/trocr-base-handwritten"),
# "model": VisionEncoderDecoderModel.from_pretrained(f"microsoft/trocr-base-handwritten"),
# "device": "cuda:0"
# },
"JorisCos/DCCRNet_Libri1Mix_enhsingle_16k": {
"model": BaseModel.from_pretrained("JorisCos/DCCRNet_Libri1Mix_enhsingle_16k"),
"device": "cuda:0"
},
# "CompVis/stable-diffusion-v1-4": {
# "model": DiffusionPipeline.from_pretrained(f"CompVis/stable-diffusion-v1-4"),
# "device": "cuda:0"
# },
# "stabilityai/stable-diffusion-2-1": {
# "model": DiffusionPipeline.from_pretrained(f"stabilityai/stable-diffusion-2-1"),
# "device": "cuda:0"
# },
# "microsoft/speecht5_tts":{
# "processor": SpeechT5Processor.from_pretrained(f"microsoft/speecht5_tts"),
# "model": SpeechT5ForTextToSpeech.from_pretrained(f"microsoft/speecht5_tts"),
# "vocoder": SpeechT5HifiGan.from_pretrained(f"microsoft/speecht5_hifigan"),
# "embeddings_dataset": load_dataset(f"Matthijs/cmu-arctic-xvectors", split="validation"),
# "device": "cuda:0"
# },
# "speechbrain/mtl-mimic-voicebank": {
# "model": WaveformEnhancement.from_hparams(source="speechbrain/mtl-mimic-voicebank", savedir="models/mtl-mimic-voicebank"),
# "device": "cuda:0"
# },
"microsoft/speecht5_vc":{
"processor": SpeechT5Processor.from_pretrained(f"{local_models}microsoft/speecht5_vc"),
"model": SpeechT5ForSpeechToSpeech.from_pretrained(f"{local_models}microsoft/speecht5_vc"),
"vocoder": SpeechT5HifiGan.from_pretrained(f"{local_models}microsoft/speecht5_hifigan"),
"embeddings_dataset": load_dataset(f"{local_models}Matthijs/cmu-arctic-xvectors", split="validation"),
"device": "cuda:0"
},
# "julien-c/wine-quality": {
# "model": joblib.load(cached_download(hf_hub_url("julien-c/wine-quality", "sklearn_model.joblib")))
# },
# "facebook/timesformer-base-finetuned-k400": {
# "processor": AutoImageProcessor.from_pretrained(f"facebook/timesformer-base-finetuned-k400"),
# "model": TimesformerForVideoClassification.from_pretrained(f"facebook/timesformer-base-finetuned-k400"),
# "device": "cuda:0"
# },
"facebook/maskformer-swin-base-coco": {
"feature_extractor": MaskFormerFeatureExtractor.from_pretrained(f"{local_models}facebook/maskformer-swin-base-coco"),
"model": MaskFormerForInstanceSegmentation.from_pretrained(f"{local_models}facebook/maskformer-swin-base-coco"),
"device": "cuda:0"
},
"Intel/dpt-hybrid-midas": {
"model": DPTForDepthEstimation.from_pretrained(f"{local_models}Intel/dpt-hybrid-midas", low_cpu_mem_usage=True),
"feature_extractor": DPTFeatureExtractor.from_pretrained(f"{local_models}Intel/dpt-hybrid-midas"),
"device": "cuda:0"
}
}
if local_deployment in ["full", "standard"]:
standard_pipes = {
# "nlpconnect/vit-gpt2-image-captioning":{
# "model": VisionEncoderDecoderModel.from_pretrained(f"{local_models}nlpconnect/vit-gpt2-image-captioning"),
# "feature_extractor": ViTImageProcessor.from_pretrained(f"{local_models}nlpconnect/vit-gpt2-image-captioning"),
# "tokenizer": AutoTokenizer.from_pretrained(f"{local_models}nlpconnect/vit-gpt2-image-captioning"),
# "device": "cuda:0"
# },
"espnet/kan-bayashi_ljspeech_vits": {
"model": Text2Speech.from_pretrained("espnet/kan-bayashi_ljspeech_vits"),
"device": "cuda:0"
},
# "lambdalabs/sd-image-variations-diffusers": {
# "model": DiffusionPipeline.from_pretrained(f"{local_models}lambdalabs/sd-image-variations-diffusers"), #torch_dtype=torch.float16
# "device": "cuda:0"
# },
"runwayml/stable-diffusion-v1-5": {
"model": DiffusionPipeline.from_pretrained(f"{local_models}runwayml/stable-diffusion-v1-5"),
"device": "cuda:0"
},
# "superb/wav2vec2-base-superb-ks": {
# "model": pipeline(task="audio-classification", model=f"superb/wav2vec2-base-superb-ks"),
# "device": "cuda:0"
# },
"openai/whisper-base": {
"model": pipeline(task="automatic-speech-recognition", model=f"{local_models}openai/whisper-base"),
"device": "cuda:0"
},
# "microsoft/speecht5_asr": {
# "model": pipeline(task="automatic-speech-recognition", model=f"{local_models}microsoft/speecht5_asr"),
# "device": "cuda:0"
# },
"Intel/dpt-large": {
"model": pipeline(task="depth-estimation", model=f"{local_models}Intel/dpt-large"),
"device": "cuda:0"
},
# "microsoft/beit-base-patch16-224-pt22k-ft22k": {
# "model": pipeline(task="image-classification", model=f"microsoft/beit-base-patch16-224-pt22k-ft22k"),
# "device": "cuda:0"
# },
"facebook/detr-resnet-50-panoptic": {
"model": pipeline(task="image-segmentation", model=f"{local_models}facebook/detr-resnet-50-panoptic"),
"device": "cuda:0"
},
"facebook/detr-resnet-101": {
"model": pipeline(task="object-detection", model=f"{local_models}facebook/detr-resnet-101"),
"device": "cuda:0"
},
# "openai/clip-vit-large-patch14": {
# "model": pipeline(task="zero-shot-image-classification", model=f"openai/clip-vit-large-patch14"),
# "device": "cuda:0"
# },
# "google/owlvit-base-patch32": {
# "model": pipeline(task="zero-shot-object-detection", model=f"{local_models}google/owlvit-base-patch32"),
# "device": "cuda:0"
# },
# "microsoft/DialoGPT-medium": {
# "model": pipeline(task="conversational", model=f"microsoft/DialoGPT-medium"),
# "device": "cuda:0"
# },
# "bert-base-uncased": {
# "model": pipeline(task="fill-mask", model=f"bert-base-uncased"),
# "device": "cuda:0"
# },
# "deepset/roberta-base-squad2": {
# "model": pipeline(task = "question-answering", model=f"deepset/roberta-base-squad2"),
# "device": "cuda:0"
# },
# "facebook/bart-large-cnn": {
# "model": pipeline(task="summarization", model=f"facebook/bart-large-cnn"),
# "device": "cuda:0"
# },
# "google/tapas-base-finetuned-wtq": {
# "model": pipeline(task="table-question-answering", model=f"google/tapas-base-finetuned-wtq"),
# "device": "cuda:0"
# },
# "distilbert-base-uncased-finetuned-sst-2-english": {
# "model": pipeline(task="text-classification", model=f"distilbert-base-uncased-finetuned-sst-2-english"),
# "device": "cuda:0"
# },
# "gpt2": {
# "model": pipeline(task="text-generation", model="gpt2"),
# "device": "cuda:0"
# },
# "mrm8488/t5-base-finetuned-question-generation-ap": {
# "model": pipeline(task="text2text-generation", model=f"mrm8488/t5-base-finetuned-question-generation-ap"),
# "device": "cuda:0"
# },
# "Jean-Baptiste/camembert-ner": {
# "model": pipeline(task="token-classification", model=f"Jean-Baptiste/camembert-ner", aggregation_strategy="simple"),
# "device": "cuda:0"
# },
# "t5-base": {
# "model": pipeline(task="translation", model=f"t5-base"),
# "device": "cuda:0"
# },
"impira/layoutlm-document-qa": {
"model": pipeline(task="document-question-answering", model=f"{local_models}impira/layoutlm-document-qa"),
"device": "cuda:0"
},
"ydshieh/vit-gpt2-coco-en": {
"model": pipeline(task="image-to-text", model=f"{local_models}ydshieh/vit-gpt2-coco-en"),
"device": "cuda:0"
},
"dandelin/vilt-b32-finetuned-vqa": {
"model": pipeline(task="visual-question-answering", model=f"{local_models}dandelin/vilt-b32-finetuned-vqa"),
"device": "cuda:0"
}
}
if local_deployment in ["full", "standard", "minimal"]:
controlnet = ControlNetModel.from_pretrained(f"{local_models}lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
controlnetpipe = StableDiffusionControlNetPipeline.from_pretrained(
f"{local_models}runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
)
hed_network = HEDdetector.from_pretrained('lllyasviel/ControlNet')
controlnet_sd_pipes = {
"openpose-control": {
"model": OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
},
"mlsd-control": {
"model": MLSDdetector.from_pretrained('lllyasviel/ControlNet')
},
"hed-control": {
"model": hed_network
},
"scribble-control": {
"model": hed_network
},
"midas-control": {
"model": MidasDetector.from_pretrained('lllyasviel/ControlNet')
},
"canny-control": {
"model": CannyDetector()
},
"lllyasviel/sd-controlnet-canny":{
"control": controlnet,
"model": controlnetpipe,
"device": "cuda:0"
},
"lllyasviel/sd-controlnet-depth":{
"control": ControlNetModel.from_pretrained(f"{local_models}lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16),
"model": controlnetpipe,
"device": "cuda:0"
},
"lllyasviel/sd-controlnet-hed":{
"control": ControlNetModel.from_pretrained(f"{local_models}lllyasviel/sd-controlnet-hed", torch_dtype=torch.float16),
"model": controlnetpipe,
"device": "cuda:0"
},
"lllyasviel/sd-controlnet-mlsd":{
"control": ControlNetModel.from_pretrained(f"{local_models}lllyasviel/sd-controlnet-mlsd", torch_dtype=torch.float16),
"model": controlnetpipe,
"device": "cuda:0"
},
"lllyasviel/sd-controlnet-openpose":{
"control": ControlNetModel.from_pretrained(f"{local_models}lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16),
"model": controlnetpipe,
"device": "cuda:0"
},
"lllyasviel/sd-controlnet-scribble":{
"control": ControlNetModel.from_pretrained(f"{local_models}lllyasviel/sd-controlnet-scribble", torch_dtype=torch.float16),
"model": controlnetpipe,
"device": "cuda:0"
},
"lllyasviel/sd-controlnet-seg":{
"control": ControlNetModel.from_pretrained(f"{local_models}lllyasviel/sd-controlnet-seg", torch_dtype=torch.float16),
"model": controlnetpipe,
"device": "cuda:0"
}
}
pipes = {**standard_pipes, **other_pipes, **controlnet_sd_pipes}
return pipes
pipes = load_pipes(local_deployment)
end = time.time()
during = end - start
print(f"[ ready ] {during}s")
def running():
return {"running": True}
def status(model_id):
disabled_models = ["microsoft/trocr-base-printed", "microsoft/trocr-base-handwritten"]
if model_id in pipes.keys() and model_id not in disabled_models:
print(f"[ check {model_id} ] success")
return {"loaded": True}
else:
print(f"[ check {model_id} ] failed")
return {"loaded": False}
def models(model_id, data):
while "using" in pipes[model_id] and pipes[model_id]["using"]:
print(f"[ inference {model_id} ] waiting")
time.sleep(0.1)
pipes[model_id]["using"] = True
print(f"[ inference {model_id} ] start")
start = time.time()
pipe = pipes[model_id]["model"]
if "device" in pipes[model_id]:
try:
pipe.to(pipes[model_id]["device"])
except:
pipe.device = torch.device(pipes[model_id]["device"])
pipe.model.to(pipes[model_id]["device"])
result = None
try:
# text to video
if model_id == "damo-vilab/text-to-video-ms-1.7b":
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
# pipe.enable_model_cpu_offload()
prompt = data["text"]
video_frames = pipe(prompt, num_inference_steps=50, num_frames=40).frames
file_name = str(uuid.uuid4())[:4]
video_path = export_to_video(video_frames, f"public/videos/{file_name}.mp4")
new_file_name = str(uuid.uuid4())[:4]
os.system(f"ffmpeg -i {video_path} -vcodec libx264 public/videos/{new_file_name}.mp4")
if os.path.exists(f"public/videos/{new_file_name}.mp4"):
result = {"path": f"/videos/{new_file_name}.mp4"}
else:
result = {"path": f"/videos/{file_name}.mp4"}
# controlnet
if model_id.startswith("lllyasviel/sd-controlnet-"):
pipe.controlnet.to('cpu')
pipe.controlnet = pipes[model_id]["control"].to(pipes[model_id]["device"])
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
control_image = load_image(data["img_url"])
# generator = torch.manual_seed(66)
out_image: Image = pipe(data["text"], num_inference_steps=20, image=control_image).images[0]
file_name = str(uuid.uuid4())[:4]
out_image.save(f"public/images/{file_name}.png")
result = {"path": f"/images/{file_name}.png"}
if model_id.endswith("-control"):
image = load_image(data["img_url"])
if "scribble" in model_id:
control = pipe(image, scribble = True)
elif "canny" in model_id:
control = pipe(image, low_threshold=100, high_threshold=200)
else:
control = pipe(image)
file_name = str(uuid.uuid4())[:4]
control.save(f"public/images/{file_name}.png")
result = {"path": f"/images/{file_name}.png"}
# image to image
if model_id == "lambdalabs/sd-image-variations-diffusers":
im = load_image(data["img_url"])
file_name = str(uuid.uuid4())[:4]
with open(f"public/images/{file_name}.png", "wb") as f:
f.write(data)
tform = transforms.Compose([
transforms.ToTensor(),
transforms.Resize(
(224, 224),
interpolation=transforms.InterpolationMode.BICUBIC,
antialias=False,
),
transforms.Normalize(
[0.48145466, 0.4578275, 0.40821073],
[0.26862954, 0.26130258, 0.27577711]),
])
inp = tform(im).to(pipes[model_id]["device"]).unsqueeze(0)
out = pipe(inp, guidance_scale=3)
out["images"][0].save(f"public/images/{file_name}.jpg")
result = {"path": f"/images/{file_name}.jpg"}
# image to text
if model_id == "Salesforce/blip-image-captioning-large":
raw_image = load_image(data["img_url"]).convert('RGB')
text = data["text"]
inputs = pipes[model_id]["processor"](raw_image, return_tensors="pt").to(pipes[model_id]["device"])
out = pipe.generate(**inputs)
caption = pipes[model_id]["processor"].decode(out[0], skip_special_tokens=True)
result = {"generated text": caption}
if model_id == "ydshieh/vit-gpt2-coco-en":
img_url = data["img_url"]
generated_text = pipe(img_url)[0]['generated_text']
result = {"generated text": generated_text}
if model_id == "nlpconnect/vit-gpt2-image-captioning":
image = load_image(data["img_url"]).convert("RGB")
pixel_values = pipes[model_id]["feature_extractor"](images=image, return_tensors="pt").pixel_values
pixel_values = pixel_values.to(pipes[model_id]["device"])
generated_ids = pipe.generate(pixel_values, **{"max_length": 200, "num_beams": 1})
generated_text = pipes[model_id]["tokenizer"].batch_decode(generated_ids, skip_special_tokens=True)[0]
result = {"generated text": generated_text}
# image to text: OCR
if model_id == "microsoft/trocr-base-printed" or model_id == "microsoft/trocr-base-handwritten":
image = load_image(data["img_url"]).convert("RGB")
pixel_values = pipes[model_id]["processor"](image, return_tensors="pt").pixel_values
pixel_values = pixel_values.to(pipes[model_id]["device"])
generated_ids = pipe.generate(pixel_values)
generated_text = pipes[model_id]["processor"].batch_decode(generated_ids, skip_special_tokens=True)[0]
result = {"generated text": generated_text}
# text to image
if model_id == "runwayml/stable-diffusion-v1-5":
file_name = str(uuid.uuid4())[:4]
text = data["text"]
out = pipe(prompt=text)
out["images"][0].save(f"public/images/{file_name}.jpg")
result = {"path": f"/images/{file_name}.jpg"}
# object detection
if model_id == "google/owlvit-base-patch32" or model_id == "facebook/detr-resnet-101":
img_url = data["img_url"]
open_types = ["cat", "couch", "person", "car", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird"]
result = pipe(img_url, candidate_labels=open_types)
# VQA
if model_id == "dandelin/vilt-b32-finetuned-vqa":
question = data["text"]
img_url = data["img_url"]
result = pipe(question=question, image=img_url)
#DQA
if model_id == "impira/layoutlm-document-qa":
question = data["text"]
img_url = data["img_url"]
result = pipe(img_url, question)
# depth-estimation
if model_id == "Intel/dpt-large":
output = pipe(data["img_url"])
image = output['depth']
name = str(uuid.uuid4())[:4]
image.save(f"public/images/{name}.jpg")
result = {"path": f"/images/{name}.jpg"}
if model_id == "Intel/dpt-hybrid-midas" and model_id == "Intel/dpt-large":
image = load_image(data["img_url"])
inputs = pipes[model_id]["feature_extractor"](images=image, return_tensors="pt")
with torch.no_grad():
outputs = pipe(**inputs)
predicted_depth = outputs.predicted_depth
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
size=image.size[::-1],
mode="bicubic",
align_corners=False,
)
output = prediction.squeeze().cpu().numpy()
formatted = (output * 255 / np.max(output)).astype("uint8")
image = Image.fromarray(formatted)
name = str(uuid.uuid4())[:4]
image.save(f"public/images/{name}.jpg")
result = {"path": f"/images/{name}.jpg"}
# TTS
if model_id == "espnet/kan-bayashi_ljspeech_vits":
text = data["text"]
wav = pipe(text)["wav"]
name = str(uuid.uuid4())[:4]
sf.write(f"public/audios/{name}.wav", wav.cpu().numpy(), pipe.fs, "PCM_16")
result = {"path": f"/audios/{name}.wav"}
if model_id == "microsoft/speecht5_tts":
text = data["text"]
inputs = pipes[model_id]["processor"](text=text, return_tensors="pt")
embeddings_dataset = pipes[model_id]["embeddings_dataset"]
speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0).to(pipes[model_id]["device"])
pipes[model_id]["vocoder"].to(pipes[model_id]["device"])
speech = pipe.generate_speech(inputs["input_ids"].to(pipes[model_id]["device"]), speaker_embeddings, vocoder=pipes[model_id]["vocoder"])
name = str(uuid.uuid4())[:4]
sf.write(f"public/audios/{name}.wav", speech.cpu().numpy(), samplerate=16000)
result = {"path": f"/audios/{name}.wav"}
# ASR
if model_id == "openai/whisper-base" or model_id == "microsoft/speecht5_asr":
audio_url = data["audio_url"]
result = { "text": pipe(audio_url)["text"]}
# audio to audio
if model_id == "JorisCos/DCCRNet_Libri1Mix_enhsingle_16k":
audio_url = data["audio_url"]
wav, sr = torchaudio.load(audio_url)
with torch.no_grad():
result_wav = pipe(wav.to(pipes[model_id]["device"]))
name = str(uuid.uuid4())[:4]
sf.write(f"public/audios/{name}.wav", result_wav.cpu().squeeze().numpy(), sr)
result = {"path": f"/audios/{name}.wav"}
if model_id == "microsoft/speecht5_vc":
audio_url = data["audio_url"]
wav, sr = torchaudio.load(audio_url)
inputs = pipes[model_id]["processor"](audio=wav, sampling_rate=sr, return_tensors="pt")
embeddings_dataset = pipes[model_id]["embeddings_dataset"]
speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
pipes[model_id]["vocoder"].to(pipes[model_id]["device"])
speech = pipe.generate_speech(inputs["input_ids"].to(pipes[model_id]["device"]), speaker_embeddings, vocoder=pipes[model_id]["vocoder"])
name = str(uuid.uuid4())[:4]
sf.write(f"public/audios/{name}.wav", speech.cpu().numpy(), samplerate=16000)
result = {"path": f"/audios/{name}.wav"}
# segmentation
if model_id == "facebook/detr-resnet-50-panoptic":
result = []
segments = pipe(data["img_url"])
image = load_image(data["img_url"])
colors = []
for i in range(len(segments)):
colors.append((random.randint(100, 255), random.randint(100, 255), random.randint(100, 255), 50))
for segment in segments:
mask = segment["mask"]
mask = mask.convert('L')
layer = Image.new('RGBA', mask.size, colors[i])
image.paste(layer, (0, 0), mask)
name = str(uuid.uuid4())[:4]
image.save(f"public/images/{name}.jpg")
result = {"path": f"/images/{name}.jpg"}
if model_id == "facebook/maskformer-swin-base-coco" or model_id == "facebook/maskformer-swin-large-ade":
image = load_image(data["img_url"])
inputs = pipes[model_id]["feature_extractor"](images=image, return_tensors="pt").to(pipes[model_id]["device"])
outputs = pipe(**inputs)
result = pipes[model_id]["feature_extractor"].post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
predicted_panoptic_map = result["segmentation"].cpu().numpy()
predicted_panoptic_map = Image.fromarray(predicted_panoptic_map.astype(np.uint8))
name = str(uuid.uuid4())[:4]
predicted_panoptic_map.save(f"public/images/{name}.jpg")
result = {"path": f"/images/{name}.jpg"}
except Exception as e:
print(e)
traceback.print_exc()
result = {"error": {"message": "Error when running the model inference."}}
if "device" in pipes[model_id]:
try:
pipe.to("cpu")
torch.cuda.empty_cache()
except:
pipe.device = torch.device("cpu")
pipe.model.to("cpu")
torch.cuda.empty_cache()
pipes[model_id]["using"] = False
if result is None:
result = {"error": {"message": "model not found"}}
end = time.time()
during = end - start
print(f"[ complete {model_id} ] {during}s")
print(f"[ result {model_id} ] {result}")
return result