Update txt2panoimg/text_to_360panorama_image_pipeline.py

txt2panoimg/text_to_360panorama_image_pipeline.py

@@ -1,19 +1,32 @@
-# Copyright © Alibaba, Inc. and its affiliates.
 import random
 from typing import Any, Dict
 
 import numpy as np
 import torch
-from basicsr.archs.rrdbnet_arch import RRDBNet
 from diffusers import (ControlNetModel, DiffusionPipeline,
                        EulerAncestralDiscreteScheduler,
                        UniPCMultistepScheduler)
 from PIL import Image
-from realesrgan import RealESRGANer
+from RealESRGAN import RealESRGAN
 
 from .pipeline_base import StableDiffusionBlendExtendPipeline
 from .pipeline_sr import StableDiffusionControlNetImg2ImgPanoPipeline
 
+class LazyRealESRGAN:
+    def __init__(self, device, scale):
+        self.device = device
+        self.scale = scale
+        self.model = None
+        self.model_path = None
+
+    def load_model(self):
+        if self.model is None:
+            self.model = RealESRGAN(self.device, scale=self.scale)
+            self.model.load_weights(self.model_path, download=False)
+
+    def predict(self, img):
+        self.load_model()
+        return self.model.predict(img)
 
 class Text2360PanoramaImagePipeline(DiffusionPipeline):
     """ Stable Diffusion for 360 Panorama Image Generation Pipeline.
@@ -40,7 +53,7 @@ class Text2360PanoramaImagePipeline(DiffusionPipeline):
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu'
                               ) if device is None else device
         if device == 'gpu':
-            device = 'cuda'
+            device = torch.device('cuda')
 
         torch_dtype = kwargs.get('torch_dtype', torch.float16)
         enable_xformers_memory_efficient_attention = kwargs.get(
@@ -60,7 +73,6 @@ class Text2360PanoramaImagePipeline(DiffusionPipeline):
                 self.pipe.enable_xformers_memory_efficient_attention()
         except Exception as e:
             print(e)
-        self.pipe.enable_model_cpu_offload()
 
         # init controlnet-sr model
         base_model_path = model + '/sr-base'
@@ -79,35 +91,15 @@ class Text2360PanoramaImagePipeline(DiffusionPipeline):
                 self.pipe_sr.enable_xformers_memory_efficient_attention()
         except Exception as e:
             print(e)
-        self.pipe_sr.enable_model_cpu_offload()
-
-        # init realesrgan model
-        sr_model = RRDBNet(
-            num_in_ch=3,
-            num_out_ch=3,
-            num_feat=64,
-            num_block=23,
-            num_grow_ch=32,
-            scale=2)
-        netscale = 2
-
+        device = torch.device("cuda")
         model_path = model + '/RealESRGAN_x2plus.pth'
-
-        dni_weight = None
-        self.upsampler = RealESRGANer(
-            scale=netscale,
-            model_path=model_path,
-            dni_weight=dni_weight,
-            model=sr_model,
-            tile=384,
-            tile_pad=20,
-            pre_pad=20,
-            half=False,
-            device=device,
-        )
+        self.upsampler = LazyRealESRGAN(device=device, scale=2)
+        self.upsampler.model_path = model_path
 
     @staticmethod
     def blend_h(a, b, blend_extent):
+        a = np.array(a)
+        b = np.array(b)
         blend_extent = min(a.shape[1], b.shape[1], blend_extent)
         for x in range(blend_extent):
             b[:, x, :] = a[:, -blend_extent
@@ -188,8 +180,8 @@ class Text2360PanoramaImagePipeline(DiffusionPipeline):
             output_img = np.array(output_img)
             output_img = np.concatenate(
                 [output_img, output_img[:, :blend_extend, :]], axis=1)
-            output_img, _ = self.upsampler.enhance(
-                output_img, outscale=outscale)
+            output_img = self.upsampler.predict(
+                output_img)
             output_img = self.blend_h(output_img, output_img,
                                       blend_extend * outscale)
             output_img = Image.fromarray(output_img[:, :w * outscale, :])