Spaces:
Paused
Paused
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 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": "cpu" | |
# }, | |
# "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": "cpu", | |
# }, | |
# "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": "cpu" | |
# }, | |
# "microsoft/trocr-base-printed": { | |
# "processor": TrOCRProcessor.from_pretrained(f"microsoft/trocr-base-printed"), | |
# "model": VisionEncoderDecoderModel.from_pretrained(f"microsoft/trocr-base-printed"), | |
# "device": "cpu" | |
# }, | |
# "microsoft/trocr-base-handwritten": { | |
# "processor": TrOCRProcessor.from_pretrained(f"microsoft/trocr-base-handwritten"), | |
# "model": VisionEncoderDecoderModel.from_pretrained(f"microsoft/trocr-base-handwritten"), | |
# "device": "cpu" | |
# }, | |
# "JorisCos/DCCRNet_Libri1Mix_enhsingle_16k": { | |
# "model": BaseModel.from_pretrained( | |
# "JorisCos/DCCRNet_Libri1Mix_enhsingle_16k" | |
# ), | |
# "device": "cpu", | |
# }, | |
# "CompVis/stable-diffusion-v1-4": { | |
# "model": DiffusionPipeline.from_pretrained(f"CompVis/stable-diffusion-v1-4"), | |
# "device": "cpu" | |
# }, | |
# "stabilityai/stable-diffusion-2-1": { | |
# "model": DiffusionPipeline.from_pretrained(f"stabilityai/stable-diffusion-2-1"), | |
# "device": "cpu" | |
# }, | |
# "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": "cpu" | |
# }, | |
# "speechbrain/mtl-mimic-voicebank": { | |
# "model": WaveformEnhancement.from_hparams(source="speechbrain/mtl-mimic-voicebank", savedir="models/mtl-mimic-voicebank"), | |
# "device": "cpu" | |
# }, | |
# "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": "cpu", | |
# }, | |
# "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": "cpu" | |
# }, | |
"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": "cpu", | |
}, | |
# "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": "cpu", | |
# }, | |
} | |
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": "cpu" | |
# }, | |
# "espnet/kan-bayashi_ljspeech_vits": { | |
# "model": Text2Speech.from_pretrained( | |
# "espnet/kan-bayashi_ljspeech_vits" | |
# ), | |
# "device": "cpu", | |
# }, | |
# "lambdalabs/sd-image-variations-diffusers": { | |
# "model": DiffusionPipeline.from_pretrained(f"{local_models}lambdalabs/sd-image-variations-diffusers"), #torch_dtype=torch.float16 | |
# "device": "cpu" | |
# }, | |
# "runwayml/stable-diffusion-v1-5": { | |
# "model": DiffusionPipeline.from_pretrained( | |
# f"{local_models}runwayml/stable-diffusion-v1-5" | |
# ), | |
# "device": "cpu", | |
# }, | |
# "superb/wav2vec2-base-superb-ks": { | |
# "model": pipeline(task="audio-classification", model=f"superb/wav2vec2-base-superb-ks"), | |
# "device": "cpu" | |
# }, | |
# "openai/whisper-base": { | |
# "model": pipeline( | |
# task="automatic-speech-recognition", | |
# model=f"{local_models}openai/whisper-base", | |
# ), | |
# "device": "cpu", | |
# }, | |
# "microsoft/speecht5_asr": { | |
# "model": pipeline(task="automatic-speech-recognition", model=f"{local_models}microsoft/speecht5_asr"), | |
# "device": "cpu" | |
# }, | |
"Intel/dpt-large": { | |
"model": pipeline( | |
task="depth-estimation", model=f"{local_models}Intel/dpt-large" | |
), | |
"device": "cpu", | |
}, | |
# "microsoft/beit-base-patch16-224-pt22k-ft22k": { | |
# "model": pipeline(task="image-classification", model=f"microsoft/beit-base-patch16-224-pt22k-ft22k"), | |
# "device": "cpu" | |
# }, | |
"facebook/detr-resnet-50-panoptic": { | |
"model": pipeline( | |
task="image-segmentation", | |
model=f"{local_models}facebook/detr-resnet-50-panoptic", | |
), | |
"device": "cpu", | |
}, | |
"facebook/detr-resnet-101": { | |
"model": pipeline( | |
task="object-detection", | |
model=f"{local_models}facebook/detr-resnet-101", | |
), | |
"device": "cpu", | |
}, | |
# "openai/clip-vit-large-patch14": { | |
# "model": pipeline(task="zero-shot-image-classification", model=f"openai/clip-vit-large-patch14"), | |
# "device": "cpu" | |
# }, | |
# "google/owlvit-base-patch32": { | |
# "model": pipeline(task="zero-shot-object-detection", model=f"{local_models}google/owlvit-base-patch32"), | |
# "device": "cpu" | |
# }, | |
# "microsoft/DialoGPT-medium": { | |
# "model": pipeline(task="conversational", model=f"microsoft/DialoGPT-medium"), | |
# "device": "cpu" | |
# }, | |
# "bert-base-uncased": { | |
# "model": pipeline(task="fill-mask", model=f"bert-base-uncased"), | |
# "device": "cpu" | |
# }, | |
# "deepset/roberta-base-squad2": { | |
# "model": pipeline(task = "question-answering", model=f"deepset/roberta-base-squad2"), | |
# "device": "cpu" | |
# }, | |
# "facebook/bart-large-cnn": { | |
# "model": pipeline(task="summarization", model=f"facebook/bart-large-cnn"), | |
# "device": "cpu" | |
# }, | |
# "google/tapas-base-finetuned-wtq": { | |
# "model": pipeline(task="table-question-answering", model=f"google/tapas-base-finetuned-wtq"), | |
# "device": "cpu" | |
# }, | |
# "distilbert-base-uncased-finetuned-sst-2-english": { | |
# "model": pipeline(task="text-classification", model=f"distilbert-base-uncased-finetuned-sst-2-english"), | |
# "device": "cpu" | |
# }, | |
# "gpt2": { | |
# "model": pipeline(task="text-generation", model="gpt2"), | |
# "device": "cpu" | |
# }, | |
# "mrm8488/t5-base-finetuned-question-generation-ap": { | |
# "model": pipeline(task="text2text-generation", model=f"mrm8488/t5-base-finetuned-question-generation-ap"), | |
# "device": "cpu" | |
# }, | |
# "Jean-Baptiste/camembert-ner": { | |
# "model": pipeline(task="token-classification", model=f"Jean-Baptiste/camembert-ner", aggregation_strategy="simple"), | |
# "device": "cpu" | |
# }, | |
# "t5-base": { | |
# "model": pipeline(task="translation", model=f"t5-base"), | |
# "device": "cpu" | |
# }, | |
# "impira/layoutlm-document-qa": { | |
# "model": pipeline(task="document-question-answering", model=f"{local_models}impira/layoutlm-document-qa"), | |
# "device": "cpu" | |
# }, | |
"ydshieh/vit-gpt2-coco-en": { | |
"model": pipeline( | |
task="image-to-text", | |
model=f"{local_models}ydshieh/vit-gpt2-coco-en", | |
), | |
"device": "cpu", | |
}, | |
# "dandelin/vilt-b32-finetuned-vqa": { | |
# "model": pipeline( | |
# task="visual-question-answering", | |
# model=f"{local_models}dandelin/vilt-b32-finetuned-vqa", | |
# ), | |
# "device": "cpu", | |
# }, | |
} | |
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") | |
pipes = {**standard_pipes, **other_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 | |