test2

app.py

@@ -77,16 +77,18 @@ def segmentation(img):
 
 def upscale(image, prompt):
     print("upscale",image,prompt)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print("device",device)
 
     # image.thumbnail((512, 512))
     # print("resize",image)
 
-    # pipeline = StableDiffusionUpscalePipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16)
-    # pipeline = pipeline.to("cuda")
-
-    pipeline = StableDiffusionLatentUpscalePipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16).to('cuda')
+    pipe = StableDiffusionUpscalePipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16)
+    # pipe = StableDiffusionLatentUpscalePipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16)
+    pipe = pipe.to(device)
+    pipe.enable_attention_slicing()
 
-    ret = pipeline(prompt=prompt, 
+    ret = pipe(prompt=prompt, 
                    image=image,
                    num_inference_steps=10,
                    guidance_scale=0)
@@ -98,12 +100,14 @@ def upscale(image, prompt):
 
 def upscale2(image, prompt):
     print("upscale2",image,prompt)
-
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipeline = LDMSuperResolutionPipeline.from_pretrained("CompVis/ldm-super-resolution-4x-openimages", torch_dtype=torch.float16)
-    pipeline = pipeline.to(device)
+    print("device",device)
 
-    upscaled_image = pipeline(image, num_inference_steps=10, eta=1).images[0]
+    pipe = LDMSuperResolutionPipeline.from_pretrained("CompVis/ldm-super-resolution-4x-openimages", torch_dtype=torch.float16)
+    pipe = pipe.to(device)
+    pipe.enable_attention_slicing()
+
+    upscaled_image = pipe(image, num_inference_steps=10, eta=1).images[0]
     return upscaled_image
 
 def convert_pil_to_cv2(image):
@@ -138,25 +142,29 @@ def pre_process(img: np.array) -> np.array:
     return img
 
 def upscale3(image):
+    print("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
+    # if img.ndim == 2:
+    #     print("upscale3","img.ndim == 2")
+    #     img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+
+    # if img.shape[2] == 4:
+    #     print("upscale3","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
+
+    # print("upscale3","img.shape[2] == 3")
+    image_output = post_process(inference(model_path, pre_process(img)))  # BGR
 
     return image_output