kinsung commited on
Commit
45191a0
Β·
verified Β·
1 Parent(s): 6a2c074
Files changed (3) hide show
  1. README.md +13 -13
  2. app.py +233 -0
  3. requirements.txt +15 -0
README.md CHANGED
@@ -1,13 +1,13 @@
1
- ---
2
- title: Ttuu
3
- emoji: πŸ“š
4
- colorFrom: blue
5
- colorTo: blue
6
- sdk: gradio
7
- sdk_version: 5.20.1
8
- app_file: app.py
9
- pinned: false
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