Update app.py

app.py

@@ -1,159 +1,233 @@
-from pathlib import Path
-import gradio as gr
-import pillow_heif
 import spaces
+
+import os
+import requests
+import time
+
 import torch
-from huggingface_hub import hf_hub_download
+
+from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+from diffusers.models import AutoencoderKL
+from diffusers.models.attention_processor import AttnProcessor2_0
+
 from PIL import Image
-from refiners.fluxion.utils import manual_seed
-from refiners.foundationals.latent_diffusion import Solver, solvers
-from enhancer import ESRGANUpscaler, ESRGANUpscalerCheckpoints
-import os
+import cv2
+import numpy as np
 
-pillow_heif.register_heif_opener()
-pillow_heif.register_avif_opener()
+from RealESRGAN import RealESRGAN
 
-CHECKPOINTS = ESRGANUpscalerCheckpoints(
-    unet=Path(
-        hf_hub_download(
-            repo_id="refiners/juggernaut.reborn.sd1_5.unet",
-            filename="model.safetensors",
-            revision="347d14c3c782c4959cc4d1bb1e336d19f7dda4d2",
-        )
-    ),
-    clip_text_encoder=Path(
-        hf_hub_download(
-            repo_id="refiners/juggernaut.reborn.sd1_5.text_encoder",
-            filename="model.safetensors",
-            revision="744ad6a5c0437ec02ad826df9f6ede102bb27481",
-        )
-    ),
-    lda=Path(
-        hf_hub_download(
-            repo_id="refiners/juggernaut.reborn.sd1_5.autoencoder",
-            filename="model.safetensors",
-            revision="3c1aae3fc3e03e4a2b7e0fa42b62ebb64f1a4c19",
-        )
-    ),
-    controlnet_tile=Path(
-        hf_hub_download(
-            repo_id="refiners/controlnet.sd1_5.tile",
-            filename="model.safetensors",
-            revision="48ced6ff8bfa873a8976fa467c3629a240643387",
+import gradio as gr
+from gradio_imageslider import ImageSlider
+
+from huggingface_hub import hf_hub_download
+
+USE_TORCH_COMPILE = False
+ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def download_models():
+    models = {
+        "MODEL": ("dantea1118/juggernaut_reborn", "juggernaut_reborn.safetensors", "models/models/Stable-diffusion"),
+        "UPSCALER_X2": ("ai-forever/Real-ESRGAN", "RealESRGAN_x2.pth", "models/upscalers/"),
+        "UPSCALER_X4": ("ai-forever/Real-ESRGAN", "RealESRGAN_x4.pth", "models/upscalers/"),
+        "NEGATIVE_1": ("philz1337x/embeddings", "verybadimagenegative_v1.3.pt", "models/embeddings"),
+        "NEGATIVE_2": ("philz1337x/embeddings", "JuggernautNegative-neg.pt", "models/embeddings"),
+        "LORA_1": ("philz1337x/loras", "SDXLrender_v2.0.safetensors", "models/Lora"),
+        "LORA_2": ("philz1337x/loras", "more_details.safetensors", "models/Lora"),
+        "CONTROLNET": ("lllyasviel/ControlNet-v1-1", "control_v11f1e_sd15_tile.pth", "models/ControlNet"),
+        "VAE": ("stabilityai/sd-vae-ft-mse-original", "vae-ft-mse-840000-ema-pruned.safetensors", "models/VAE"),
+    }
+
+    for model, (repo_id, filename, local_dir) in models.items():
+        hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir)
+
+download_models()
+
+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...")
+            self.pipe = self.setup_pipeline()
+            print(f"Moving pipeline to device: {device}")
+            self.pipe.to(device)
+            if USE_TORCH_COMPILE:
+                print("Compiling the model...")
+                self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
+
+    @timer_func
+    def setup_pipeline(self):
+        print("Setting up the pipeline...")
+        controlnet = ControlNetModel.from_single_file(
+            "models/ControlNet/control_v11f1e_sd15_tile.pth", torch_dtype=torch.float16
         )
-    ),
-    esrgan=Path(
-        hf_hub_download(
-            repo_id="philz1337x/upscaler",
-            filename="4x-UltraSharp.pth",
-            revision="011deacac8270114eb7d2eeff4fe6fa9a837be70",
+        safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+        model_path = "models/models/Stable-diffusion/juggernaut_reborn.safetensors"
+        pipe = StableDiffusionControlNetImg2ImgPipeline.from_single_file(
+            model_path,
+            controlnet=controlnet,
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+            safety_checker=safety_checker
         )
-    ),
-    negative_embedding=Path(
-        hf_hub_download(
-            repo_id="philz1337x/embeddings",
-            filename="JuggernautNegative-neg.pt",
-            revision="203caa7e9cc2bc225031a4021f6ab1ded283454a",
+        vae = AutoencoderKL.from_single_file(
+            "models/VAE/vae-ft-mse-840000-ema-pruned.safetensors",
+            torch_dtype=torch.float16
         )
-    ),
-    negative_embedding_key="string_to_param.*",
-    loras={
-        "more_details": Path(
-            hf_hub_download(
-                repo_id="philz1337x/loras",
-                filename="more_details.safetensors",
-                revision="a3802c0280c0d00c2ab18d37454a8744c44e474e",
-            )
-        ),
-        "sdxl_render": Path(
-            hf_hub_download(
-                repo_id="philz1337x/loras",
-                filename="SDXLrender_v2.0.safetensors",
-                revision="a3802c0280c0d00c2ab18d37454a8744c44e474e",
-            )
-        ),
-    },
-)
-
-# initialize the enhancer, on the cpu
-DEVICE_CPU = torch.device("cpu")
-DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
-enhancer = ESRGANUpscaler(checkpoints=CHECKPOINTS, device=DEVICE_CPU, dtype=DTYPE)
-
-# "move" the enhancer to the gpu, this is handled by Zero GPU
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-enhancer.to(device=DEVICE, dtype=DTYPE)
-
+        pipe.vae = vae
+        pipe.load_textual_inversion("models/embeddings/verybadimagenegative_v1.3.pt")
+        pipe.load_textual_inversion("models/embeddings/JuggernautNegative-neg.pt")
+        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")
+        pipe.fuse_lora(lora_scale=1.)
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4)
+        return pipe
+
+    def __call__(self, *args, **kwargs):
+        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()
+lazy_pipe.load()
+
+def prepare_image(input_image, resolution, hdr):
+    condition_image = resize_and_upscale(input_image, resolution)
+    condition_image = create_hdr_effect(condition_image, hdr)
+    return condition_image
 
 @spaces.GPU
-def process(
-    input_image: Image.Image,
-    denoise_strength: float,
-) -> tuple[Image.Image, str]:
-    # 하드코딩된 옵션
+@timer_func
+def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
+    print("Starting image processing...")
+    torch.cuda.empty_cache()
+    
+    condition_image = prepare_image(input_image, resolution, hdr)
+    
     prompt = "masterpiece, best quality, highres"
-    negative_prompt = "worst quality, low quality, normal quality"
-    seed = 42
-    upscale_factor = 2
-    controlnet_scale = 0.6
-    controlnet_decay = 1.0
-    condition_scale = 6
-    tile_width = 112
-    tile_height = 144
-    num_inference_steps = 18
-    solver = "DDIM"
-
-    manual_seed(seed)
-
-    solver_type: type[Solver] = getattr(solvers, solver)
-
-    enhanced_image = enhancer.upscale(
-        image=input_image,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        upscale_factor=upscale_factor,
-        controlnet_scale=controlnet_scale,
-        controlnet_scale_decay=controlnet_decay,
-        condition_scale=condition_scale,
-        tile_size=(tile_height, tile_width),
-        denoise_strength=denoise_strength,
-        num_inference_steps=num_inference_steps,
-        loras_scale={"more_details": 0.5, "sdxl_render": 1.0},
-        solver_type=solver_type,
-    )
-
-    # 임시로 이미지를 저장하여 다운로드 가능하게 처리
-    output_path = "/tmp/enhanced_image.jpg"
-    enhanced_image.save(output_path, "JPEG")
-
-    return enhanced_image, output_path
-
+    negative_prompt = "low quality, normal quality, ugly, blurry, blur, lowres, bad anatomy, bad hands, cropped, worst quality, verybadimagenegative_v1.3, JuggernautNegative-neg"
+    
+    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),
+    }
+    
+    print("Running inference...")
+    result = lazy_pipe(**options).images[0]
+    print("Image processing completed successfully")
+    
+    # Convert input_image and result to numpy arrays
+    input_array = np.array(input_image)
+    result_array = np.array(result)
+    
+    return [input_array, result_array]
+
+title = """<h1 align="center">Image Upscaler with Tile Controlnet</h1>
+<p align="center">The main ideas come from</p>
+<p><center>
+<a href="https://github.com/philz1337x/clarity-upscaler" target="_blank">[philz1337x]</a>
+<a href="https://github.com/BatouResearch/controlnet-tile-upscale" target="_blank">[Pau-Lozano]</a>
+</center></p>
+"""
 
 with gr.Blocks() as demo:
-
+    gr.HTML(title)
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(type="pil", label="Input Image")
-            run_button = gr.Button(value="Enhance Image")
+            run_button = gr.Button("Enhance Image")
         with gr.Column():
-            output_image = gr.Image(label="Enhanced Image")
-            download_file = gr.File(label="Download Enhanced Image")
-
-    # 강도 조절을 위한 라디오 버튼 추가 (3단계)
-    denoise_strength = gr.Radio(
-        choices=[
-            ("약하게", 0.3),
-            ("중간", 0.5),
-            ("강하게", 0.8)
+            output_slider = ImageSlider(label="Before / After", type="numpy")
+    with gr.Accordion("Advanced Options", open=False):
+        resolution = gr.Slider(minimum=256, maximum=2048, value=512, step=256, label="Resolution")
+        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.4, step=0.01, 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")
+
+    run_button.click(fn=gradio_process_image, 
+                     inputs=[input_image, resolution, num_inference_steps, strength, hdr, guidance_scale],
+                     outputs=output_slider)
+
+    # Add examples with all required inputs
+    gr.Examples(
+        examples=[
+            ["image1.jpg", 512, 20, 0.4, 0, 3],
+            ["image2.png", 512, 20, 0.4, 0, 3],
+            ["image3.png", 512, 20, 0.4, 0, 3],
         ],
-        label="강도 조절",
-        value=0.3  # 기본값: 3단계
-    )
-
-    run_button.click(
-        fn=process,
-        inputs=[input_image, denoise_strength],
-        outputs=[output_image, download_file],
+        inputs=[input_image, resolution, num_inference_steps, strength, hdr, guidance_scale],
+        outputs=output_slider,
+        fn=gradio_process_image,
+        cache_examples=True,
     )
 
-demo.launch(share=True)
+demo.launch(share=True)