up 3

app.py

@@ -20,7 +20,8 @@ from diffusers import DiffusionPipeline
 from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionPipeline
 from diffusers import StableDiffusionUpscalePipeline
 from diffusers import LDMSuperResolutionPipeline
-
+import cv2
+import onnxruntime
 
 def removeFurniture(input_img1,
             input_img2,
@@ -105,7 +106,59 @@ def upscale2(image, prompt):
     upscaled_image = pipeline(image, num_inference_steps=10, eta=1).images[0]
     return upscaled_image
 
+def convert_pil_to_cv2(image):
+    # pil_image = image.convert("RGB")
+    open_cv_image = np.array(image)
+    # RGB to BGR
+    open_cv_image = open_cv_image[:, :, ::-1].copy()
+    return open_cv_image
+
+def inference(model_path: str, img_array: np.array) -> np.array:
+    options = onnxruntime.SessionOptions()
+    options.intra_op_num_threads = 1
+    options.inter_op_num_threads = 1
+    ort_session = onnxruntime.InferenceSession(model_path, options)
+    ort_inputs = {ort_session.get_inputs()[0].name: img_array}
+    ort_outs = ort_session.run(None, ort_inputs)
+
+    return ort_outs[0]
+
+def post_process(img: np.array) -> np.array:
+    # 1, C, H, W -> C, H, W
+    img = np.squeeze(img)
+    # C, H, W -> H, W, C
+    img = np.transpose(img, (1, 2, 0))[:, :, ::-1].astype(np.uint8)
+    return img
+
+def pre_process(img: np.array) -> np.array:
+    # H, W, C -> C, H, W
+    img = np.transpose(img[:, :, 0:3], (2, 0, 1))
+    # C, H, W -> 1, C, H, W
+    img = np.expand_dims(img, axis=0).astype(np.float32)
+    return img
+
+def upscale3(image):
+    model_path = f"up_models/modelx4.ort"
+    img = convert_pil_to_cv2(image)
+    
+    if img.ndim == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+
+    if img.shape[2] == 4:
+        alpha = img[:, :, 3]  # GRAY
+        alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)  # BGR
+        alpha_output = post_process(inference(model_path, pre_process(alpha)))  # BGR
+        alpha_output = cv2.cvtColor(alpha_output, cv2.COLOR_BGR2GRAY)  # GRAY
+
+        img = img[:, :, 0:3]  # BGR
+        image_output = post_process(inference(model_path, pre_process(img)))  # BGR
+        image_output = cv2.cvtColor(image_output, cv2.COLOR_BGR2BGRA)  # BGRA
+        image_output[:, :, 3] = alpha_output
 
+    elif img.shape[2] == 3:
+        image_output = post_process(inference(model_path, pre_process(img)))  # BGR
+
+    return image_output
 
 
 with gr.Blocks() as app:    
@@ -141,6 +194,9 @@ with gr.Blocks() as app:
         with gr.Column():
             gr.Button("Upscale2").click(upscale2, inputs=[gr.Image(type="pil"),gr.Textbox(label="prompt",value="empty room")], outputs=gr.Image())
 
+        with gr.Column():
+            gr.Button("Upscale3").click(upscale3, inputs=[gr.Image(type="pil")], outputs=gr.Image())
+
 
 app.launch(debug=True,share=True)
 

requirements.txt

@@ -9,3 +9,4 @@ python-docx
 triton
 altair<5
 gradio
+onnxruntime==1.12.0