|
import spaces |
|
import torch |
|
|
|
|
|
|
|
def my_arange(*args, **kwargs): |
|
return torch.arange(*args, **kwargs) |
|
|
|
|
|
torch.arange = my_arange |
|
|
|
from pathlib import Path |
|
|
|
import gradio as gr |
|
from gradio_imageslider import ImageSlider |
|
from huggingface_hub import hf_hub_download |
|
from PIL import Image |
|
from refiners.fluxion.utils import manual_seed |
|
from refiners.foundationals.latent_diffusion import Solver, solvers |
|
|
|
from enhancer import ESRGANUpscaler, ESRGANUpscalerCheckpoints |
|
|
|
TITLE = """ |
|
<h1 align="center">Image Enhancer, implemented using refiners</h1> |
|
|
|
<p> |
|
<center> |
|
<a style="font-size: 1.25rem;" href="https://blog.finegrain.ai/posts/reproducing-clarity-upscaler/" target="_blank">[blog post]</a> |
|
<a style="font-size: 1.25rem;" href="https://github.com/finegrain-ai/refiners" target="_blank">[refiners]</a> |
|
<a style="font-size: 1.25rem;" href="https://github.com/philz1337x/clarity-upscaler" target="_blank">[clarity-upscaler]</a> |
|
<a style="font-size: 1.25rem;" href="https://finegrain.ai/" target="_blank">[finegrain]</a> |
|
</center> |
|
</p> |
|
""" |
|
|
|
CHECKPOINTS = ESRGANUpscalerCheckpoints( |
|
unet=Path( |
|
hf_hub_download( |
|
repo_id="refiners/juggernaut.reborn", |
|
filename="unet.safetensors", |
|
revision="948510aaf4c8e8e9b32b5a7c25736422253f7b93", |
|
) |
|
), |
|
clip_text_encoder=Path( |
|
hf_hub_download( |
|
repo_id="refiners/juggernaut.reborn", |
|
filename="text_encoder.safetensors", |
|
revision="948510aaf4c8e8e9b32b5a7c25736422253f7b93", |
|
) |
|
), |
|
lda=Path( |
|
hf_hub_download( |
|
repo_id="refiners/juggernaut.reborn", |
|
filename="autoencoder.safetensors", |
|
revision="948510aaf4c8e8e9b32b5a7c25736422253f7b93", |
|
) |
|
), |
|
controlnet_tile=Path( |
|
hf_hub_download( |
|
repo_id="refiners/controlnet.sd15.tile", |
|
filename="model.safetensors", |
|
revision="48ced6ff8bfa873a8976fa467c3629a240643387", |
|
) |
|
), |
|
esrgan=Path( |
|
hf_hub_download( |
|
repo_id="philz1337x/upscaler", |
|
filename="4x-UltraSharp.pth", |
|
revision="011deacac8270114eb7d2eeff4fe6fa9a837be70", |
|
) |
|
), |
|
negative_embedding=Path( |
|
hf_hub_download( |
|
repo_id="philz1337x/embeddings", |
|
filename="JuggernautNegative-neg.pt", |
|
revision="203caa7e9cc2bc225031a4021f6ab1ded283454a", |
|
) |
|
), |
|
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", |
|
) |
|
), |
|
}, |
|
) |
|
|
|
LORA_SCALES = { |
|
"more_details": 0.5, |
|
"sdxl_render": 1.0, |
|
} |
|
|
|
|
|
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) |
|
|
|
|
|
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
enhancer.to(device=DEVICE, dtype=DTYPE) |
|
|
|
|
|
@spaces.GPU |
|
def process( |
|
input_image: Image.Image, |
|
prompt: str = "masterpiece, best quality, highres", |
|
negative_prompt: str = "worst quality, low quality, normal quality", |
|
seed: int = 42, |
|
upscale_factor: int = 2, |
|
controlnet_scale: float = 0.6, |
|
controlnet_decay: float = 1.0, |
|
condition_scale: int = 6, |
|
tile_width: int = 112, |
|
tile_height: int = 144, |
|
denoise_strength: float = 0.35, |
|
num_inference_steps: int = 18, |
|
solver: str = "DDIM", |
|
) -> tuple[Image.Image, Image.Image]: |
|
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=LORA_SCALES, |
|
solver_type=solver_type, |
|
) |
|
|
|
return (input_image, enhanced_image) |
|
|
|
|
|
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.ClearButton(components=None, value="Enhance Image") |
|
with gr.Column(): |
|
output_slider = ImageSlider(label="Before / After") |
|
run_button.add(output_slider) |
|
|
|
with gr.Accordion("Advanced Options", open=False): |
|
prompt = gr.Textbox( |
|
label="Prompt", |
|
placeholder="masterpiece, best quality, highres", |
|
) |
|
negative_prompt = gr.Textbox( |
|
label="Negative Prompt", |
|
placeholder="worst quality, low quality, normal quality", |
|
) |
|
seed = gr.Slider( |
|
minimum=0, |
|
maximum=10_000, |
|
value=42, |
|
step=1, |
|
label="Seed", |
|
) |
|
upscale_factor = gr.Slider( |
|
minimum=1, |
|
maximum=4, |
|
value=2, |
|
step=0.2, |
|
label="Upscale Factor", |
|
) |
|
controlnet_scale = gr.Slider( |
|
minimum=0, |
|
maximum=1.5, |
|
value=0.6, |
|
step=0.1, |
|
label="ControlNet Scale", |
|
) |
|
controlnet_decay = gr.Slider( |
|
minimum=0.5, |
|
maximum=1, |
|
value=1.0, |
|
step=0.025, |
|
label="ControlNet Scale Decay", |
|
) |
|
condition_scale = gr.Slider( |
|
minimum=2, |
|
maximum=20, |
|
value=6, |
|
step=1, |
|
label="Condition Scale", |
|
) |
|
tile_width = gr.Slider( |
|
minimum=64, |
|
maximum=200, |
|
value=112, |
|
step=1, |
|
label="Latent Tile Width", |
|
) |
|
tile_height = gr.Slider( |
|
minimum=64, |
|
maximum=200, |
|
value=144, |
|
step=1, |
|
label="Latent Tile Height", |
|
) |
|
denoise_strength = gr.Slider( |
|
minimum=0, |
|
maximum=1, |
|
value=0.35, |
|
step=0.1, |
|
label="Denoise Strength", |
|
) |
|
num_inference_steps = gr.Slider( |
|
minimum=1, |
|
maximum=30, |
|
value=18, |
|
step=1, |
|
label="Number of Inference Steps", |
|
) |
|
solver = gr.Radio( |
|
choices=["DDIM", "DPMSolver"], |
|
value="DDIM", |
|
label="Solver", |
|
) |
|
|
|
run_button.click( |
|
fn=process, |
|
inputs=[ |
|
input_image, |
|
prompt, |
|
negative_prompt, |
|
seed, |
|
upscale_factor, |
|
controlnet_scale, |
|
controlnet_decay, |
|
condition_scale, |
|
tile_width, |
|
tile_height, |
|
denoise_strength, |
|
num_inference_steps, |
|
solver, |
|
], |
|
outputs=output_slider, |
|
) |
|
|
|
gr.Examples( |
|
examples=[ |
|
"examples/kara-eads-L7EwHkq1B2s-unsplash.jpg", |
|
"examples/clarity_bird.webp", |
|
"examples/edgar-infocus-gJH8AqpiSEU-unsplash.jpg", |
|
"examples/jeremy-wallace-_XjW3oN8UOE-unsplash.jpg", |
|
"examples/karina-vorozheeva-rW-I87aPY5Y-unsplash.jpg", |
|
"examples/karographix-photography-hIaOPjYCEj4-unsplash.jpg", |
|
"examples/melissa-walker-horn-gtDYwUIr9Vg-unsplash.jpg", |
|
"examples/ryoji-iwata-X53e51WfjlE-unsplash.jpg", |
|
"examples/tadeusz-lakota-jggQZkITXng-unsplash.jpg", |
|
], |
|
inputs=[input_image], |
|
outputs=output_slider, |
|
fn=process, |
|
cache_examples="lazy", |
|
run_on_click=False, |
|
) |
|
|
|
demo.launch(share=False) |
|
|