upload
Browse files- README.md +13 -13
- app.py +233 -0
- requirements.txt +15 -0
README.md
CHANGED
|
@@ -1,13 +1,13 @@
|
|
| 1 |
-
---
|
| 2 |
-
title:
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo: blue
|
| 6 |
-
sdk: gradio
|
| 7 |
-
sdk_version:
|
| 8 |
-
app_file: app.py
|
| 9 |
-
pinned:
|
| 10 |
-
license: apache-2.0
|
| 11 |
-
---
|
| 12 |
-
|
| 13 |
-
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: Tile Upscaler
|
| 3 |
+
emoji: 🚀
|
| 4 |
+
colorFrom: purple
|
| 5 |
+
colorTo: blue
|
| 6 |
+
sdk: gradio
|
| 7 |
+
sdk_version: 4.37.2
|
| 8 |
+
app_file: app.py
|
| 9 |
+
pinned: true
|
| 10 |
+
license: apache-2.0
|
| 11 |
+
---
|
| 12 |
+
|
| 13 |
+
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
app.py
ADDED
|
@@ -0,0 +1,233 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import spaces
|
| 2 |
+
|
| 3 |
+
import os
|
| 4 |
+
import requests
|
| 5 |
+
import time
|
| 6 |
+
|
| 7 |
+
import torch
|
| 8 |
+
|
| 9 |
+
from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler
|
| 10 |
+
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
|
| 11 |
+
from diffusers.models import AutoencoderKL
|
| 12 |
+
from diffusers.models.attention_processor import AttnProcessor2_0
|
| 13 |
+
|
| 14 |
+
from PIL import Image
|
| 15 |
+
import cv2
|
| 16 |
+
import numpy as np
|
| 17 |
+
|
| 18 |
+
from RealESRGAN import RealESRGAN
|
| 19 |
+
|
| 20 |
+
import gradio as gr
|
| 21 |
+
from gradio_imageslider import ImageSlider
|
| 22 |
+
|
| 23 |
+
from huggingface_hub import hf_hub_download
|
| 24 |
+
|
| 25 |
+
USE_TORCH_COMPILE = False
|
| 26 |
+
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
|
| 27 |
+
|
| 28 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 29 |
+
|
| 30 |
+
def download_models():
|
| 31 |
+
models = {
|
| 32 |
+
"MODEL": ("dantea1118/juggernaut_reborn", "juggernaut_reborn.safetensors", "models/models/Stable-diffusion"),
|
| 33 |
+
"UPSCALER_X2": ("ai-forever/Real-ESRGAN", "RealESRGAN_x2.pth", "models/upscalers/"),
|
| 34 |
+
"UPSCALER_X4": ("ai-forever/Real-ESRGAN", "RealESRGAN_x4.pth", "models/upscalers/"),
|
| 35 |
+
"NEGATIVE_1": ("philz1337x/embeddings", "verybadimagenegative_v1.3.pt", "models/embeddings"),
|
| 36 |
+
"NEGATIVE_2": ("philz1337x/embeddings", "JuggernautNegative-neg.pt", "models/embeddings"),
|
| 37 |
+
"LORA_1": ("philz1337x/loras", "SDXLrender_v2.0.safetensors", "models/Lora"),
|
| 38 |
+
"LORA_2": ("philz1337x/loras", "more_details.safetensors", "models/Lora"),
|
| 39 |
+
"CONTROLNET": ("lllyasviel/ControlNet-v1-1", "control_v11f1e_sd15_tile.pth", "models/ControlNet"),
|
| 40 |
+
"VAE": ("stabilityai/sd-vae-ft-mse-original", "vae-ft-mse-840000-ema-pruned.safetensors", "models/VAE"),
|
| 41 |
+
}
|
| 42 |
+
|
| 43 |
+
for model, (repo_id, filename, local_dir) in models.items():
|
| 44 |
+
hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir)
|
| 45 |
+
|
| 46 |
+
download_models()
|
| 47 |
+
|
| 48 |
+
def timer_func(func):
|
| 49 |
+
def wrapper(*args, **kwargs):
|
| 50 |
+
start_time = time.time()
|
| 51 |
+
result = func(*args, **kwargs)
|
| 52 |
+
end_time = time.time()
|
| 53 |
+
print(f"{func.__name__} took {end_time - start_time:.2f} seconds")
|
| 54 |
+
return result
|
| 55 |
+
return wrapper
|
| 56 |
+
|
| 57 |
+
class LazyLoadPipeline:
|
| 58 |
+
def __init__(self):
|
| 59 |
+
self.pipe = None
|
| 60 |
+
|
| 61 |
+
@timer_func
|
| 62 |
+
def load(self):
|
| 63 |
+
if self.pipe is None:
|
| 64 |
+
print("Starting to load the pipeline...")
|
| 65 |
+
self.pipe = self.setup_pipeline()
|
| 66 |
+
print(f"Moving pipeline to device: {device}")
|
| 67 |
+
self.pipe.to(device)
|
| 68 |
+
if USE_TORCH_COMPILE:
|
| 69 |
+
print("Compiling the model...")
|
| 70 |
+
self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
|
| 71 |
+
|
| 72 |
+
@timer_func
|
| 73 |
+
def setup_pipeline(self):
|
| 74 |
+
print("Setting up the pipeline...")
|
| 75 |
+
controlnet = ControlNetModel.from_single_file(
|
| 76 |
+
"models/ControlNet/control_v11f1e_sd15_tile.pth", torch_dtype=torch.float16
|
| 77 |
+
)
|
| 78 |
+
safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
|
| 79 |
+
model_path = "models/models/Stable-diffusion/juggernaut_reborn.safetensors"
|
| 80 |
+
pipe = StableDiffusionControlNetImg2ImgPipeline.from_single_file(
|
| 81 |
+
model_path,
|
| 82 |
+
controlnet=controlnet,
|
| 83 |
+
torch_dtype=torch.float16,
|
| 84 |
+
use_safetensors=True,
|
| 85 |
+
safety_checker=safety_checker
|
| 86 |
+
)
|
| 87 |
+
vae = AutoencoderKL.from_single_file(
|
| 88 |
+
"models/VAE/vae-ft-mse-840000-ema-pruned.safetensors",
|
| 89 |
+
torch_dtype=torch.float16
|
| 90 |
+
)
|
| 91 |
+
pipe.vae = vae
|
| 92 |
+
pipe.load_textual_inversion("models/embeddings/verybadimagenegative_v1.3.pt")
|
| 93 |
+
pipe.load_textual_inversion("models/embeddings/JuggernautNegative-neg.pt")
|
| 94 |
+
pipe.load_lora_weights("models/Lora/SDXLrender_v2.0.safetensors")
|
| 95 |
+
pipe.fuse_lora(lora_scale=0.5)
|
| 96 |
+
pipe.load_lora_weights("models/Lora/more_details.safetensors")
|
| 97 |
+
pipe.fuse_lora(lora_scale=1.)
|
| 98 |
+
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
|
| 99 |
+
pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4)
|
| 100 |
+
return pipe
|
| 101 |
+
|
| 102 |
+
def __call__(self, *args, **kwargs):
|
| 103 |
+
return self.pipe(*args, **kwargs)
|
| 104 |
+
|
| 105 |
+
class LazyRealESRGAN:
|
| 106 |
+
def __init__(self, device, scale):
|
| 107 |
+
self.device = device
|
| 108 |
+
self.scale = scale
|
| 109 |
+
self.model = None
|
| 110 |
+
|
| 111 |
+
def load_model(self):
|
| 112 |
+
if self.model is None:
|
| 113 |
+
self.model = RealESRGAN(self.device, scale=self.scale)
|
| 114 |
+
self.model.load_weights(f'models/upscalers/RealESRGAN_x{self.scale}.pth', download=False)
|
| 115 |
+
def predict(self, img):
|
| 116 |
+
self.load_model()
|
| 117 |
+
return self.model.predict(img)
|
| 118 |
+
|
| 119 |
+
lazy_realesrgan_x2 = LazyRealESRGAN(device, scale=2)
|
| 120 |
+
lazy_realesrgan_x4 = LazyRealESRGAN(device, scale=4)
|
| 121 |
+
|
| 122 |
+
@timer_func
|
| 123 |
+
def resize_and_upscale(input_image, resolution):
|
| 124 |
+
scale = 2 if resolution <= 2048 else 4
|
| 125 |
+
input_image = input_image.convert("RGB")
|
| 126 |
+
W, H = input_image.size
|
| 127 |
+
k = float(resolution) / min(H, W)
|
| 128 |
+
H = int(round(H * k / 64.0)) * 64
|
| 129 |
+
W = int(round(W * k / 64.0)) * 64
|
| 130 |
+
img = input_image.resize((W, H), resample=Image.LANCZOS)
|
| 131 |
+
if scale == 2:
|
| 132 |
+
img = lazy_realesrgan_x2.predict(img)
|
| 133 |
+
else:
|
| 134 |
+
img = lazy_realesrgan_x4.predict(img)
|
| 135 |
+
return img
|
| 136 |
+
|
| 137 |
+
@timer_func
|
| 138 |
+
def create_hdr_effect(original_image, hdr):
|
| 139 |
+
if hdr == 0:
|
| 140 |
+
return original_image
|
| 141 |
+
cv_original = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
|
| 142 |
+
factors = [1.0 - 0.9 * hdr, 1.0 - 0.7 * hdr, 1.0 - 0.45 * hdr,
|
| 143 |
+
1.0 - 0.25 * hdr, 1.0, 1.0 + 0.2 * hdr,
|
| 144 |
+
1.0 + 0.4 * hdr, 1.0 + 0.6 * hdr, 1.0 + 0.8 * hdr]
|
| 145 |
+
images = [cv2.convertScaleAbs(cv_original, alpha=factor) for factor in factors]
|
| 146 |
+
merge_mertens = cv2.createMergeMertens()
|
| 147 |
+
hdr_image = merge_mertens.process(images)
|
| 148 |
+
hdr_image_8bit = np.clip(hdr_image * 255, 0, 255).astype('uint8')
|
| 149 |
+
return Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))
|
| 150 |
+
|
| 151 |
+
lazy_pipe = LazyLoadPipeline()
|
| 152 |
+
lazy_pipe.load()
|
| 153 |
+
|
| 154 |
+
def prepare_image(input_image, resolution, hdr):
|
| 155 |
+
condition_image = resize_and_upscale(input_image, resolution)
|
| 156 |
+
condition_image = create_hdr_effect(condition_image, hdr)
|
| 157 |
+
return condition_image
|
| 158 |
+
|
| 159 |
+
@spaces.GPU
|
| 160 |
+
@timer_func
|
| 161 |
+
def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
|
| 162 |
+
print("Starting image processing...")
|
| 163 |
+
torch.cuda.empty_cache()
|
| 164 |
+
|
| 165 |
+
condition_image = prepare_image(input_image, resolution, hdr)
|
| 166 |
+
|
| 167 |
+
prompt = "masterpiece, best quality, highres"
|
| 168 |
+
negative_prompt = "low quality, normal quality, ugly, blurry, blur, lowres, bad anatomy, bad hands, cropped, worst quality, verybadimagenegative_v1.3, JuggernautNegative-neg"
|
| 169 |
+
|
| 170 |
+
options = {
|
| 171 |
+
"prompt": prompt,
|
| 172 |
+
"negative_prompt": negative_prompt,
|
| 173 |
+
"image": condition_image,
|
| 174 |
+
"control_image": condition_image,
|
| 175 |
+
"width": condition_image.size[0],
|
| 176 |
+
"height": condition_image.size[1],
|
| 177 |
+
"strength": strength,
|
| 178 |
+
"num_inference_steps": num_inference_steps,
|
| 179 |
+
"guidance_scale": guidance_scale,
|
| 180 |
+
"generator": torch.Generator(device=device).manual_seed(0),
|
| 181 |
+
}
|
| 182 |
+
|
| 183 |
+
print("Running inference...")
|
| 184 |
+
result = lazy_pipe(**options).images[0]
|
| 185 |
+
print("Image processing completed successfully")
|
| 186 |
+
|
| 187 |
+
# Convert input_image and result to numpy arrays
|
| 188 |
+
input_array = np.array(input_image)
|
| 189 |
+
result_array = np.array(result)
|
| 190 |
+
|
| 191 |
+
return [input_array, result_array]
|
| 192 |
+
|
| 193 |
+
title = """<h1 align="center">Image Upscaler with Tile Controlnet</h1>
|
| 194 |
+
<p align="center">The main ideas come from</p>
|
| 195 |
+
<p><center>
|
| 196 |
+
<a href="https://github.com/philz1337x/clarity-upscaler" target="_blank">[philz1337x]</a>
|
| 197 |
+
<a href="https://github.com/BatouResearch/controlnet-tile-upscale" target="_blank">[Pau-Lozano]</a>
|
| 198 |
+
</center></p>
|
| 199 |
+
"""
|
| 200 |
+
|
| 201 |
+
with gr.Blocks() as demo:
|
| 202 |
+
gr.HTML(title)
|
| 203 |
+
with gr.Row():
|
| 204 |
+
with gr.Column():
|
| 205 |
+
input_image = gr.Image(type="pil", label="Input Image")
|
| 206 |
+
run_button = gr.Button("Enhance Image")
|
| 207 |
+
with gr.Column():
|
| 208 |
+
output_slider = ImageSlider(label="Before / After", type="numpy")
|
| 209 |
+
with gr.Accordion("Advanced Options", open=False):
|
| 210 |
+
resolution = gr.Slider(minimum=256, maximum=2048, value=512, step=256, label="Resolution")
|
| 211 |
+
num_inference_steps = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Number of Inference Steps")
|
| 212 |
+
strength = gr.Slider(minimum=0, maximum=1, value=0.4, step=0.01, label="Strength")
|
| 213 |
+
hdr = gr.Slider(minimum=0, maximum=1, value=0, step=0.1, label="HDR Effect")
|
| 214 |
+
guidance_scale = gr.Slider(minimum=0, maximum=20, value=3, step=0.5, label="Guidance Scale")
|
| 215 |
+
|
| 216 |
+
run_button.click(fn=gradio_process_image,
|
| 217 |
+
inputs=[input_image, resolution, num_inference_steps, strength, hdr, guidance_scale],
|
| 218 |
+
outputs=output_slider)
|
| 219 |
+
|
| 220 |
+
# Add examples with all required inputs
|
| 221 |
+
gr.Examples(
|
| 222 |
+
examples=[
|
| 223 |
+
["image1.jpg", 512, 20, 0.4, 0, 3],
|
| 224 |
+
["image2.png", 512, 20, 0.4, 0, 3],
|
| 225 |
+
["image3.png", 512, 20, 0.4, 0, 3],
|
| 226 |
+
],
|
| 227 |
+
inputs=[input_image, resolution, num_inference_steps, strength, hdr, guidance_scale],
|
| 228 |
+
outputs=output_slider,
|
| 229 |
+
fn=gradio_process_image,
|
| 230 |
+
cache_examples=True,
|
| 231 |
+
)
|
| 232 |
+
|
| 233 |
+
demo.launch(share=True)
|
requirements.txt
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
git+https://github.com/doevent/Real-ESRGAN.git
|
| 2 |
+
opencv-python
|
| 3 |
+
spaces
|
| 4 |
+
diffusers
|
| 5 |
+
torch
|
| 6 |
+
torchvision
|
| 7 |
+
pipeline
|
| 8 |
+
transformers
|
| 9 |
+
accelerate
|
| 10 |
+
safetensors
|
| 11 |
+
spaces
|
| 12 |
+
peft
|
| 13 |
+
gradio
|
| 14 |
+
pillow
|
| 15 |
+
gradio-imageslider
|