Update app.py

app.py

@@ -151,6 +151,25 @@ pipe.unet.set_attn_processor(AttnProcessor2_0())
 # Enable FreeU
 pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4)
 
+class LazyRealESRGAN:
+    def __init__(self, device, scale):
+        self.device = device
+        self.scale = scale
+        self.model = None
+
+    def load_model(self):
+        if self.model is None:
+            self.model = RealESRGAN(self.device, scale=self.scale)
+            self.model.load_weights(f'models/upscalers/RealESRGAN_x{self.scale}.pth', download=False)
+
+    def predict(self, img):
+        self.load_model()
+        return self.model.predict(img)
+
+# Initialize the lazy models
+lazy_realesrgan_x2 = LazyRealESRGAN(device, scale=2)
+lazy_realesrgan_x4 = LazyRealESRGAN(device, scale=4)
+
 def resize_and_upscale(input_image, resolution):
     scale = 2
     if resolution == 2048:
@@ -174,6 +193,10 @@ def resize_and_upscale(input_image, resolution):
     model = RealESRGAN(device, scale=scale)
     model.load_weights(f'models/upscalers/RealESRGAN_x{scale}.pth', download=False)
     img = model.predict(img)
+    if scale == 2:
+        img = lazy_realesrgan_x2.predict(img)
+    else:
+        img = lazy_realesrgan_x4.predict(img)
     return img
 
 def calculate_brightness_factors(hdr_intensity):