From ce66402345cfd8d5e873edffb63a58156a276357 Mon Sep 17 00:00:00 2001
From: "baiguan.yt" <baiguan.yt@alibaba-inc.com>
Date: Thu, 4 Aug 2022 17:26:02 +0800
Subject: [PATCH] [to #42322933]fix some issues         Link:
 https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/9643850

---
 .../item_embedding.py                         |  4 ----
 .../cv/face_image_generation_pipeline.py      |  5 +++-
 .../cv/image_colorization_pipeline.py         | 23 ++++++++-----------
 .../cv/image_super_resolution_pipeline.py     |  5 +++-
 4 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/modelscope/models/cv/product_retrieval_embedding/item_embedding.py b/modelscope/models/cv/product_retrieval_embedding/item_embedding.py
index b01031d5..0444596c 100644
--- a/modelscope/models/cv/product_retrieval_embedding/item_embedding.py
+++ b/modelscope/models/cv/product_retrieval_embedding/item_embedding.py
@@ -1,9 +1,5 @@
-import os
-import time
-
 import cv2
 import numpy as np
-import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
diff --git a/modelscope/pipelines/cv/face_image_generation_pipeline.py b/modelscope/pipelines/cv/face_image_generation_pipeline.py
index 31c97b30..405c9a4b 100644
--- a/modelscope/pipelines/cv/face_image_generation_pipeline.py
+++ b/modelscope/pipelines/cv/face_image_generation_pipeline.py
@@ -29,7 +29,10 @@ class FaceImageGenerationPipeline(Pipeline):
             model: model id on modelscope hub.
         """
         super().__init__(model=model, **kwargs)
-        self.device = 'cpu'
+        if torch.cuda.is_available():
+            self.device = torch.device('cuda')
+        else:
+            self.device = torch.device('cpu')
         self.size = 1024
         self.latent = 512
         self.n_mlp = 8
diff --git a/modelscope/pipelines/cv/image_colorization_pipeline.py b/modelscope/pipelines/cv/image_colorization_pipeline.py
index 838ccab5..ff0e27ff 100644
--- a/modelscope/pipelines/cv/image_colorization_pipeline.py
+++ b/modelscope/pipelines/cv/image_colorization_pipeline.py
@@ -12,7 +12,7 @@ from modelscope.models.cv.image_colorization import (DynamicUnetDeep,
 from modelscope.outputs import OutputKeys
 from modelscope.pipelines.base import Input, Pipeline
 from modelscope.pipelines.builder import PIPELINES
-from modelscope.preprocessors import load_image
+from modelscope.preprocessors import LoadImage, load_image
 from modelscope.utils.constant import ModelFile, Tasks
 from modelscope.utils.logger import get_logger
 
@@ -31,7 +31,13 @@ class ImageColorizationPipeline(Pipeline):
         """
         super().__init__(model=model, **kwargs)
         self.cut = 8
-        self.size = 1024 if self.device_name == 'cpu' else 512
+        self.size = 512
+        if torch.cuda.is_available():
+            self.device = torch.device('cuda')
+        else:
+            self.device = torch.device('cpu')
+            self.size = 1024
+
         self.orig_img = None
         self.model_type = 'stable'
         self.norm = transforms.Compose([
@@ -82,18 +88,7 @@ class ImageColorizationPipeline(Pipeline):
         logger.info('load model done')
 
     def preprocess(self, input: Input) -> Dict[str, Any]:
-        if isinstance(input, str):
-            img = load_image(input).convert('LA').convert('RGB')
-        elif isinstance(input, Image.Image):
-            img = input.convert('LA').convert('RGB')
-        elif isinstance(input, np.ndarray):
-            if len(input.shape) == 2:
-                input = cv2.cvtColor(input, cv2.COLOR_GRAY2BGR)
-            img = input[:, :, ::-1]  # in rgb order
-            img = PIL.Image.fromarray(img).convert('LA').convert('RGB')
-        else:
-            raise TypeError(f'input should be either str, PIL.Image,'
-                            f' np.array, but got {type(input)}')
+        img = LoadImage.convert_to_img(input).convert('LA').convert('RGB')
 
         self.wide, self.height = img.size
         if self.wide * self.height > self.size * self.size:
diff --git a/modelscope/pipelines/cv/image_super_resolution_pipeline.py b/modelscope/pipelines/cv/image_super_resolution_pipeline.py
index 6464fe69..27f8bfb2 100644
--- a/modelscope/pipelines/cv/image_super_resolution_pipeline.py
+++ b/modelscope/pipelines/cv/image_super_resolution_pipeline.py
@@ -28,7 +28,10 @@ class ImageSuperResolutionPipeline(Pipeline):
             model: model id on modelscope hub.
         """
         super().__init__(model=model, **kwargs)
-        self.device = 'cpu'
+        if torch.cuda.is_available():
+            self.device = torch.device('cuda')
+        else:
+            self.device = torch.device('cpu')
         self.num_feat = 64
         self.num_block = 23
         self.scale = 4