Create models_server.py

models_server.py

@@ -0,0 +1,258 @@
+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