Update app.py

app.py

@@ -1,8 +1,7 @@
-import spaces
 import os
 import requests
 import torch
-from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler
+from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler
 from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 from diffusers.models import AutoencoderKL
 from PIL import Image
@@ -11,15 +10,13 @@ import cv2
 import numpy as np
 from diffusers.models.attention_processor import AttnProcessor2_0
 import gradio as gr
+import spaces
 
-USE_TORCH_COMPILE = 0
-ENABLE_CPU_OFFLOAD = 0
+USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
+ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
-# Set up the device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-
-# Function to download files (from the example)
 def download_file(url, folder_path, filename):
     if not os.path.exists(folder_path):
         os.makedirs(folder_path)
@@ -37,7 +34,6 @@ def download_file(url, folder_path, filename):
         else:
             print(f"Error downloading the file. Status code: {response.status_code}")
 
-# Download necessary models and files
 def download_models():
     models = {
         "MODEL": ("https://huggingface.co/dantea1118/juggernaut_reborn/resolve/main/juggernaut_reborn.safetensors?download=true", "models/models/Stable-diffusion", "juggernaut_reborn.safetensors"),
@@ -56,9 +52,6 @@ def download_models():
 
 download_models()
 
-
-
-
 class LazyRealESRGAN:
     def __init__(self, device, scale):
         self.device = device
@@ -74,72 +67,49 @@ class LazyRealESRGAN:
         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:
-        init_w = 1024
-    elif resolution == 2560:
-        init_w = 1280
-    elif resolution == 3072:
-        init_w = 1536
-    else:
-        init_w = 1024
-        scale = 4
-
+    scale = 2 if resolution <= 2048 else 4
     input_image = input_image.convert("RGB")
     W, H = input_image.size
-    k = float(init_w) / min(H, W)
-    H *= k
-    W *= k
-    H = int(round(H / 64.0)) * 64
-    W = int(round(W / 64.0)) * 64
+    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)
-    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):
-    factors = [1.0] * 9
-    if hdr_intensity > 0:
-        factors = [1.0 - 0.9 * hdr_intensity, 1.0 - 0.7 * hdr_intensity, 1.0 - 0.45 * hdr_intensity,
-                   1.0 - 0.25 * hdr_intensity, 1.0, 1.0 + 0.2 * hdr_intensity,
-                   1.0 + 0.4 * hdr_intensity, 1.0 + 0.6 * hdr_intensity, 1.0 + 0.8 * hdr_intensity]
-    return factors
-
-def pil_to_cv(pil_image):
-    return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
-
-def adjust_brightness(cv_image, factor):
-    hsv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2HSV)
-    h, s, v = cv2.split(hsv_image)
-    v = np.clip(v * factor, 0, 255).astype('uint8')
-    adjusted_hsv = cv2.merge([h, s, v])
-    return cv2.cvtColor(adjusted_hsv, cv2.COLOR_HSV2BGR)
-
 def create_hdr_effect(original_image, hdr):
-    cv_original = pil_to_cv(original_image)
-    brightness_factors = calculate_brightness_factors(hdr)
-    images = [adjust_brightness(cv_original, factor) for factor in brightness_factors]
-
+    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')
-    hdr_image_pil = Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))
-
-    return hdr_image_pil
+    return Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))
 
-class ImageProcessor:
+class LazyLoadPipeline:
     def __init__(self):
-        self.pipe = self.setup_pipeline()
+        self.pipe = None
+
+    def load(self):
+        if self.pipe is None:
+            self.pipe = self.setup_pipeline()
+            if ENABLE_CPU_OFFLOAD:
+                self.pipe.enable_model_cpu_offload()
+            else:
+                self.pipe.to(device)
+            if USE_TORCH_COMPILE:
+                self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
 
     def setup_pipeline(self):
         controlnet = ControlNetModel.from_single_file(
@@ -172,35 +142,38 @@ class ImageProcessor:
 
         return pipe
 
-    def process_image(self, input_image, prompt, negative_prompt, resolution=2048, num_inference_steps=50, guidance_scale=3, strength=0.35, hdr=0):
-        condition_image = resize_and_upscale(input_image, resolution)
-        condition_image = create_hdr_effect(condition_image, hdr)
+    def __call__(self, *args, **kwargs):
+        self.load()
+        return self.pipe(*args, **kwargs)
 
-        result = self.pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=condition_image,
-            control_image=condition_image,
-            width=condition_image.size[0],
-            height=condition_image.size[1],
-            strength=strength,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            generator=torch.manual_seed(0),
-        ).images[0]
-
-        return result
-
-
-image_processor = ImageProcessor()
+lazy_pipe = LazyLoadPipeline()
 
 @spaces.GPU
 def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
-    image_processor.pipe = image_processor.pipe.to(device)
-    image_processor.pipe.unet.set_attn_processor(AttnProcessor2_0())
+    torch.cuda.empty_cache()
+    lazy_pipe.load()
+    lazy_pipe.pipe.unet.set_attn_processor(AttnProcessor2_0())
+    
+    condition_image = resize_and_upscale(input_image, resolution)
+    condition_image = create_hdr_effect(condition_image, hdr)
+    
     prompt = "masterpiece, best quality, highres"
     negative_prompt = "low quality, normal quality, ugly, blurry, blur, lowres, bad anatomy, bad hands, cropped, worst quality, verybadimagenegative_v1.3, JuggernautNegative-neg"
-    result = image_processor.process_image(input_image, prompt, negative_prompt, resolution, num_inference_steps, guidance_scale, strength, hdr)
+    
+    options = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "image": condition_image,
+        "control_image": condition_image,
+        "width": condition_image.size[0],
+        "height": condition_image.size[1],
+        "strength": strength,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "generator": torch.Generator(device=device).manual_seed(0),
+    }
+    
+    result = lazy_pipe(**options).images[0]
     return result
 
 # Gradio interface
@@ -214,7 +187,7 @@ with gr.Blocks() as demo:
             output_image = gr.Image(type="pil", label="Enhanced Image")
     with gr.Accordion("Advanced Options", open=False):
         resolution = gr.Slider(minimum=512, maximum=2048, value=1024, step=64, label="Resolution")
-        num_inference_steps = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Number of Inference Steps")
+        num_inference_steps = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Number of Inference Steps")
         strength = gr.Slider(minimum=0, maximum=1, value=0.35, step=0.05, label="Strength")
         hdr = gr.Slider(minimum=0, maximum=1, value=0, step=0.1, label="HDR Effect")
         guidance_scale = gr.Slider(minimum=0, maximum=20, value=3, step=0.5, label="Guidance Scale")