Change to a Tile Upscaler

Changes a models to Tile upscaler.

.tool-versions

@@ -0,0 +1 @@
+python 3.13.1

README.md

@@ -1,12 +1,13 @@
 ---
-title: Hackathon 2025
-emoji: 🖼
+title: Tile Upscaler
+emoji: 🚀
 colorFrom: purple
-colorTo: red
+colorTo: blue
 sdk: gradio
-sdk_version: 5.25.2
+sdk_version: 4.44.1
 app_file: app.py
-pinned: false
+pinned: true
+license: apache-2.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,294 +1,233 @@
-import gradio as gr
-import numpy as np
-
-# import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import (
-    # StableDiffusionControlNetImg2ImgPipeline,
-    ControlNetModel,
-    ControlNetUnionModel,
-    StableDiffusionXLControlNetPipeline,
-)
-import torch
+import spaces
 
+import os
 import requests
-from fastapi import FastAPI, HTTPException
-from PIL import Image
-from controlnet_aux import CannyDetector
-
-from diffusers import AutoencoderKL
-from diffusers import (
-    EulerAncestralDiscreteScheduler,
-    StableDiffusionXLControlNetUnionPipeline,
-)
-import cv2
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-# model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-model_repo_id = "runwayml/stable-diffusion-v1-5"
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-# controlnet = ControlNetModel.from_pretrained(
-#     "lllyasviel/sd-controlnet-canny", torch_dtype=torch.float32
-# )
-
-# controlnet = ControlNetModel.from_pretrained(
-#     "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16
-# )
-
-# pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-# pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-#     model_repo_id,
-#     controlnet=controlnet,
-#     torch_dtype=torch_dtype,
-# ).to(device)
-# pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-#     "stabilityai/stable-diffusion-xl-base-1.0",
-#     controlnet=controlnet,
-#     torch_dtype=torch.float16,
-#     variant="fp16",
-#     use_safetensors=True,
-# ).to(device)
-# # pipe = pipe.to(device)
-# canny = CannyDetector()
-
-
-eulera_scheduler = EulerAncestralDiscreteScheduler.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler"
-)
-
-# when test with other base model, you need to change the vae also.
-vae = AutoencoderKL.from_pretrained(
-    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
-)
-
-controlnet_model = ControlNetUnionModel.from_pretrained(
-    "xinsir/controlnet-union-sdxl-1.0", torch_dtype=torch.float16, use_safetensors=True
-)
-
-# controlnet_union_model = ControlNetUnionModel([controlnet_model])
-
-pipe = StableDiffusionXLControlNetUnionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    controlnet=controlnet_model,
-    vae=vae,
-    torch_dtype=torch.float16,
-    scheduler=eulera_scheduler,
-    control_mode=[0],
-)
-
-pipe = pipe.to(device)
-
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    image_url,
-    # negative_prompt,
-    # seed,
-    # randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    # if randomize_seed:
-    #     seed = random.randint(0, MAX_SEED)
-
-    # generator = torch.Generator().manual_seed(seed)
-
-    # image = pipe(
-    #     prompt=prompt,
-    #     negative_prompt=negative_prompt,
-    #     guidance_scale=guidance_scale,
-    #     num_inference_steps=num_inference_steps,
-    #     width=width,
-    #     height=height,
-    #     generator=generator,
-    # ).images[0]
-
-    # return image, seed
-
-    width = int(width)
-    height = int(height)
+import time
 
-    try:
-        resp = requests.get(image_url)
-        resp.raise_for_status()
-    except Exception as e:
-        raise HTTPException(400, f"Could not download image: {e}")
+import subprocess
+subprocess.run("pip install git+https://github.com/inference-sh/Real-ESRGAN.git --no-deps", shell=True)
 
-    # img = Image.open(io.BytesIO(resp.content)).convert("RGB")
-    img = Image.open(requests.get(image_url, stream=True).raw).convert("RGB")
-    # img = img.resize((req.width, req.height))
-    # img = img.resize((width, height))
-
-    # control_net_image = canny(img).resize((width, height))
-
-    img_np = np.array(img)
-
-    controlnet_img = cv2.resize(img_np, (width, height))
-
-    controlnet_img = cv2.Canny(controlnet_img, 100, 200)
-    controlnet_img = HWC3(controlnet_img)
-    controlnet_img = Image.fromarray(controlnet_img)
-
-    prompt = (
-        "redraw the logo from scratch, clean sharp vector-style, "
-        # + STYLE_PROMPTS[req.style_preset]
-    )
-
-    output = pipe(
-        prompt=prompt,
-        negative_prompt=NEGATIVE,
-        # image=img,
-        control_image=controlnet_img,
-        # strength=req.strength,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        height=height,
-        width=width,
-    ).images[0]
+import torch
 
-    return output
+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
+import cv2
+import numpy as np
 
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
+from RealESRGAN import RealESRGAN
 
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
+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
+        )
+        # 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
+        )
+        vae = AutoencoderKL.from_single_file(
+            "models/VAE/vae-ft-mse-840000-ema-pruned.safetensors",
+            torch_dtype=torch.float16
+        )
+        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
+@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 = "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>
 """
 
-NEGATIVE = "blurry, distorted, messy, gradients, background noise"
-WIDTH = 512
-HEIGHT = 512
-
-
-def HWC3(x):
-    assert x.dtype == np.uint8
-    if x.ndim == 2:
-        x = x[:, :, None]
-    assert x.ndim == 3
-    H, W, C = x.shape
-    assert C == 1 or C == 3 or C == 4
-    if C == 3:
-        return x
-    if C == 1:
-        return np.concatenate([x, x, x], axis=2)
-    if C == 4:
-        color = x[:, :, 0:3].astype(np.float32)
-        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
-        y = color * alpha + 255.0 * (1.0 - alpha)
-        y = y.clip(0, 255).astype(np.uint8)
-        return y
-
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            image_url = gr.Text(
-                label="Image URL",
-                show_label=False,
-                # max_lines=1,
-                placeholder="Provide a image URL",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Label(
-                label="Negative prompts",
-                # max_lines=1,
-                value=NEGATIVE,
-                visible=True,
-            )
-
-            # seed = gr.Slider(
-            #     label="Seed",
-            #     minimum=0,
-            #     maximum=MAX_SEED,
-            #     step=1,
-            #     value=0,
-            # )
-
-            # randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Label(
-                    label="Width",
-                    value=WIDTH,
-                    # minimum=256,
-                    # maximum=MAX_IMAGE_SIZE,
-                    # step=32,
-                    # value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Label(
-                    label="Height",
-                    value=HEIGHT,
-                    # minimum=256,
-                    # maximum=MAX_IMAGE_SIZE,
-                    # step=32,
-                    # value=1024,  # Replace with defaults that work for your model
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=8.5,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=25,  # Replace with defaults that work for your model
-                )
-
-        # gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, image_url.submit],
-        fn=infer,
-        inputs=[
-            image_url,
-            # negative_prompt,
-            # seed,
-            # randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[
-            result,
-            # seed,
+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("Enhance Image")
+        with gr.Column():
+            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=[
         ],
+        inputs=[input_image, resolution, num_inference_steps, strength, hdr, guidance_scale],
+        outputs=output_slider,
+        fn=gradio_process_image,
+        cache_examples=True,
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch(share=True)

requirements.txt

@@ -1,13 +1,15 @@
-accelerate
+opencv-python
+spaces
 diffusers
-invisible_watermark
-torch
-transformers
-# xformers
-fastapi
-uvicorn
-pydantic
-requests
-Pillow
-controlnet-aux
-gradio
+torch==2.4.0
+torchvision
+pipeline
+transformers<=4.49.0
+accelerate
+safetensors
+spaces
+peft
+gradio
+pillow
+gradio-imageslider
+pydantic==2.10.6