Update app.py

app.py

@@ -52,61 +52,25 @@ def download_models():
 
 download_models()
 
-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)
-
-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 else 4
-    input_image = input_image.convert("RGB")
-    W, H = input_image.size
-    k = float(resolution) / min(H, W)
-    H = int(round(H * k / 64.0)) * 64
-    W = int(round(W * k / 64.0)) * 64
-    img = input_image.resize((W, H), resample=Image.LANCZOS)
-    if scale == 2:
-        img = lazy_realesrgan_x2.predict(img)
-    else:
-        img = lazy_realesrgan_x4.predict(img)
-    return img
-
-def create_hdr_effect(original_image, hdr):
-    if hdr == 0:
-        return original_image
-    cv_original = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
-    factors = [1.0 - 0.9 * hdr, 1.0 - 0.7 * hdr, 1.0 - 0.45 * hdr,
-               1.0 - 0.25 * hdr, 1.0, 1.0 + 0.2 * hdr,
-               1.0 + 0.4 * hdr, 1.0 + 0.6 * hdr, 1.0 + 0.8 * hdr]
-    images = [cv2.convertScaleAbs(cv_original, alpha=factor) for factor in factors]
-    merge_mertens = cv2.createMergeMertens()
-    hdr_image = merge_mertens.process(images)
-    hdr_image_8bit = np.clip(hdr_image * 255, 0, 255).astype('uint8')
-    return Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))
-
 import time
 
+def timer_func(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        print(f"{func.__name__} took {end_time - start_time:.2f} seconds")
+        return result
+    return wrapper
+
 class LazyLoadPipeline:
     def __init__(self):
         self.pipe = None
 
+    @timer_func
     def load(self):
         if self.pipe is None:
             print("Starting to load the pipeline...")
-            start_time = time.time()
             try:
                 self.pipe = self.setup_pipeline()
                 if ENABLE_CPU_OFFLOAD:
@@ -118,18 +82,24 @@ class LazyLoadPipeline:
                 if USE_TORCH_COMPILE:
                     print("Compiling the model...")
                     self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
-                print(f"Pipeline loaded successfully in {time.time() - start_time:.2f} seconds")
             except Exception as e:
                 print(f"Error loading pipeline: {str(e)}")
                 raise
 
+    @timer_func
     def setup_pipeline(self):
         print("Setting up the pipeline...")
+        start_time = time.time()
         controlnet = ControlNetModel.from_single_file(
             "models/ControlNet/control_v11f1e_sd15_tile.pth", torch_dtype=torch.float16
         )
+        print(f"ControlNet loaded in {time.time() - start_time:.2f} seconds")
+
+        start_time = time.time()
         safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+        print(f"Safety checker loaded in {time.time() - start_time:.2f} seconds")
 
+        start_time = time.time()
         model_path = "models/models/Stable-diffusion/juggernaut_reborn.safetensors"
         pipe = StableDiffusionControlNetImg2ImgPipeline.from_single_file(
             model_path,
@@ -138,21 +108,32 @@ class LazyLoadPipeline:
             use_safetensors=True,
             safety_checker=safety_checker
         )
+        print(f"Main pipeline loaded in {time.time() - start_time:.2f} seconds")
 
+        start_time = time.time()
         vae = AutoencoderKL.from_single_file(
             "models/VAE/vae-ft-mse-840000-ema-pruned.safetensors",
             torch_dtype=torch.float16
         )
         pipe.vae = vae
+        print(f"VAE loaded in {time.time() - start_time:.2f} seconds")
 
         print("Loading textual inversions and LoRA weights...")
+        start_time = time.time()
         pipe.load_textual_inversion("models/embeddings/verybadimagenegative_v1.3.pt")
         pipe.load_textual_inversion("models/embeddings/JuggernautNegative-neg.pt")
+        print(f"Textual inversions loaded in {time.time() - start_time:.2f} seconds")
+
+        start_time = time.time()
         pipe.load_lora_weights("models/Lora/SDXLrender_v2.0.safetensors")
         pipe.fuse_lora(lora_scale=0.5)
         pipe.load_lora_weights("models/Lora/more_details.safetensors")
+        print(f"LoRA weights loaded in {time.time() - start_time:.2f} seconds")
+
+        start_time = time.time()
         pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
         pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4)
+        print(f"Scheduler and FreeU set up in {time.time() - start_time:.2f} seconds")
 
         return pipe
 
@@ -160,9 +141,57 @@ class LazyLoadPipeline:
         self.load()
         return self.pipe(*args, **kwargs)
 
+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)
+
+lazy_realesrgan_x2 = LazyRealESRGAN(device, scale=2)
+lazy_realesrgan_x4 = LazyRealESRGAN(device, scale=4)
+
+@timer_func
+def resize_and_upscale(input_image, resolution):
+    scale = 2 if resolution <= 2048 else 4
+    input_image = input_image.convert("RGB")
+    W, H = input_image.size
+    k = float(resolution) / min(H, W)
+    H = int(round(H * k / 64.0)) * 64
+    W = int(round(W * k / 64.0)) * 64
+    img = input_image.resize((W, H), resample=Image.LANCZOS)
+    if scale == 2:
+        img = lazy_realesrgan_x2.predict(img)
+    else:
+        img = lazy_realesrgan_x4.predict(img)
+    return img
+    
+@timer_func
+def create_hdr_effect(original_image, hdr):
+    if hdr == 0:
+        return original_image
+    cv_original = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
+    factors = [1.0 - 0.9 * hdr, 1.0 - 0.7 * hdr, 1.0 - 0.45 * hdr,
+               1.0 - 0.25 * hdr, 1.0, 1.0 + 0.2 * hdr,
+               1.0 + 0.4 * hdr, 1.0 + 0.6 * hdr, 1.0 + 0.8 * hdr]
+    images = [cv2.convertScaleAbs(cv_original, alpha=factor) for factor in factors]
+    merge_mertens = cv2.createMergeMertens()
+    hdr_image = merge_mertens.process(images)
+    hdr_image_8bit = np.clip(hdr_image * 255, 0, 255).astype('uint8')
+    return Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))
+
 lazy_pipe = LazyLoadPipeline()
 
 @spaces.GPU
+@timer_func
 def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
     print("Starting image processing...")
     torch.cuda.empty_cache()