app.py

import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import torch
import numpy as np
import gradio as gr
import spaces
from PIL import Image

from diffusers import DDPMScheduler
from pipelines.lcm_single_step_scheduler import LCMSingleStepScheduler

from module.ip_adapter.utils import load_adapter_to_pipe
from pipelines.sdxl_SAKBIR import SAKBIRPipeline

def resize_img(input_image, max_side=1280, min_side=1024, size=None, 
               pad_to_max_side=False, mode=Image.BILINEAR, base_pixel_number=64):

    w, h = input_image.size
    if size is not None:
        w_resize_new, h_resize_new = size
    else:
        # ratio = min_side / min(h, w)
        # w, h = round(ratio*w), round(ratio*h)
        ratio = max_side / max(h, w)
        input_image = input_image.resize([round(ratio*w), round(ratio*h)], mode)
        w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
        h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
    input_image = input_image.resize([w_resize_new, h_resize_new], mode)

    if pad_to_max_side:
        res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
        offset_x = (max_side - w_resize_new) // 2
        offset_y = (max_side - h_resize_new) // 2
        res[offset_y:offset_y+h_resize_new, offset_x:offset_x+w_resize_new] = np.array(input_image)
        input_image = Image.fromarray(res)
    return input_image

from huggingface_hub import hf_hub_download

hf_hub_download(repo_id="SunderAli17/SAKBIR", filename="models/adapter.pt", local_dir=".")
hf_hub_download(repo_id="SunderAli17/SAKBIR", filename="models/aggregator.pt", local_dir=".")
hf_hub_download(repo_id="SunderAli17/SAKBIR", filename="models/previewer_lora_weights.bin", local_dir=".")

SAKBIR_path = f'./models'

device = "cuda" if torch.cuda.is_available() else "cpu"
sdxl_repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
dinov2_repo_id = "facebook/dinov2-large"
lcm_repo_id = "latent-consistency/lcm-lora-sdxl"

if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

# Load pretrained models.
print("Initializing pipeline...")
pipe = SAKBIRPipeline.from_pretrained(
    sdxl_repo_id,
    torch_dtype=torch_dtype,
)

# Image prompt projector.
print("Loading LQ-Adapter...")
load_adapter_to_pipe(
    pipe,
    f"{SAKBIR_path}/adapter.pt",
    dinov2_repo_id,
)

# Prepare previewer
lora_alpha = pipe.prepare_previewers(SAKBIR_path)
print(f"use lora alpha {lora_alpha}")
lora_alpha = pipe.prepare_previewers(lcm_repo_id, use_lcm=True)
print(f"use lora alpha {lora_alpha}")
pipe.to(device=device, dtype=torch_dtype)
pipe.scheduler = DDPMScheduler.from_pretrained(sdxl_repo_id, subfolder="scheduler")
lcm_scheduler = LCMSingleStepScheduler.from_config(pipe.scheduler.config)

# Load weights.
print("Loading checkpoint...")
aggregator_state_dict = torch.load(
    f"{SAKBIR_path}/aggregator.pt",
    map_location="cpu"
)
pipe.aggregator.load_state_dict(aggregator_state_dict, strict=True)
pipe.aggregator.to(device=device, dtype=torch_dtype)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

PROMPT = "Photorealistic, highly detailed, hyper detailed photo - realistic maximum detail, 32k, \
ultra HD, extreme meticulous detailing, skin pore detailing, \
hyper sharpness, perfect without deformations, \
taken using a Canon EOS R camera, Cinematic, High Contrast, Color Grading. "

NEG_PROMPT = "blurry, out of focus, unclear, depth of field, over-smooth, \
sketch, oil painting, cartoon, CG Style, 3D render, unreal engine, \
dirty, messy, worst quality, low quality, frames, painting, illustration, drawing, art, \
watermark, signature, jpeg artifacts, deformed, lowres"

def unpack_pipe_out(preview_row, index):
    return preview_row[index][0]

def dynamic_preview_slider(sampling_steps):
    print(sampling_steps)
    return gr.Slider(label="Restoration Previews", value=sampling_steps-1, minimum=0, maximum=sampling_steps-1, step=1)

def dynamic_guidance_slider(sampling_steps):
    return gr.Slider(label="Start Free Rendering", value=sampling_steps, minimum=0, maximum=sampling_steps, step=1)

def show_final_preview(preview_row):
    return preview_row[-1][0]

@spaces.GPU(duration=70) #[uncomment to use ZeroGPU]
@torch.no_grad()
def SAKBIR_restore(
    lq, prompt="", steps=30, cfg_scale=7.0, guidance_end=1.0,
    creative_restoration=False, seed=3407, height=1024, width=1024, preview_start=0.0, progress=gr.Progress(track_tqdm=True)):
    if creative_restoration:
        if "lcm" not in pipe.unet.active_adapters():
            pipe.unet.set_adapter('lcm')
    else:
        if "previewer" not in pipe.unet.active_adapters():
            pipe.unet.set_adapter('previewer')

    if isinstance(guidance_end, int):
        guidance_end = guidance_end / steps
    elif guidance_end > 1.0:
        guidance_end = guidance_end / steps
    if isinstance(preview_start, int):
        preview_start = preview_start / steps
    elif preview_start > 1.0:
        preview_start = preview_start / steps

    w, h = lq.size
    if w == h :
        lq = [resize_img(lq.convert("RGB"), size=(width, height))]
    else:
        lq = [resize_img(lq.convert("RGB"), size=None)]
   
    generator = torch.Generator(device=device).manual_seed(seed)
    timesteps = [
        i * (1000//steps) + pipe.scheduler.config.steps_offset for i in range(0, steps)
    ]
    timesteps = timesteps[::-1]

    prompt = PROMPT if len(prompt)==0 else prompt
    neg_prompt = NEG_PROMPT

    out = pipe(
        prompt=[prompt]*len(lq),
        image=lq,
        num_inference_steps=steps,
        generator=generator,
        timesteps=timesteps,
        negative_prompt=[neg_prompt]*len(lq),
        guidance_scale=cfg_scale,
        control_guidance_end=guidance_end,
        preview_start=preview_start,
        previewer_scheduler=lcm_scheduler,
        return_dict=False,
        save_preview_row=True,
    )
    for i, preview_img in enumerate(out[1]):
        preview_img.append(f"preview_{i}")
    return out[0][0], out[1]

examples = [
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    "An astronaut riding a green horse",
    "A delicious ceviche cheesecake slice",
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks() as SAK:
    gr.Markdown(
    """
    # SAKBIR: Blind Image Restoration using Generative Networks.
    ### **SAKBIR can restore your degraded image as well as modify it through the imaginative prompts being provided by the user. You can refer to the advance usage details for more information.**
    ## Basic usage: Image Restoration
    1. Upload the image that you want to restore;
    2. Optionally, tune the `Steps` `CFG Scale` parameters. Typically higher steps lead to better results, but less than 50 is recommended for efficiency;
    3. Click `SAKBIR!`.
    """)
    with gr.Row():
        with gr.Column():
            lq_img = gr.Image(label="Low-quality image", type="pil")      
            
            with gr.Row():
                steps = gr.Number(label="Steps", value=30, step=1)
                cfg_scale = gr.Number(label="CFG Scale", value=7.0, step=0.1)
            
            with gr.Row():
                height = gr.Number(label="Height", value=1024, step=1, visible=False)
                width = gr.Number(label="Width", value=1024, step=1, visible=False)
                seed = gr.Number(label="Seed", value=42, step=1)
            # guidance_start = gr.Slider(label="Guidance Start", value=1.0, minimum=0.0, maximum=1.0, step=0.05)
            guidance_end = gr.Slider(label="Start Free Rendering", value=30, minimum=0, maximum=30, step=1)
            preview_start = gr.Slider(label="Preview Start", value=0, minimum=0, maximum=30, step=1)
            prompt = gr.Textbox(label="Restoration prompts (Optional)", placeholder="")
            mode = gr.Checkbox(label="Creative Restoration", value=False)
    
            with gr.Row():
                restore_btn = gr.Button("SAKBIR processing!")
                clear_btn = gr.ClearButton()
            gr.Examples(
                    examples = ["example_images/lady.png", "example_images/man.png", "example_images/dog.png", "example_images/panda.png", "example_images/sculpture.png", "example_images/cottage.png", "example_images/Naruto.png", "example_images/Konan.png"],
                    inputs = [lq_img]
                )
        with gr.Column():
            output = gr.Image(label="Image restored using SAKBIR", type="pil")
            index = gr.Slider(label="Preview Restoration", value=29, minimum=0, maximum=29, step=1)
            preview = gr.Image(label="Preview", type="pil")
       
    pipe_out = gr.Gallery(visible=False)
    clear_btn.add([lq_img, output, preview])
    restore_btn.click(
        SAKBIR_restore, inputs=[
            lq_img, prompt, steps, cfg_scale, guidance_end,
            mode, seed, height, width, preview_start,
        ],
        outputs=[output, pipe_out], api_name="SAKBIR"
    )
    steps.change(dynamic_guidance_slider, inputs=steps, outputs=guidance_end)
    output.change(dynamic_preview_slider, inputs=steps, outputs=index)
    index.release(unpack_pipe_out, inputs=[pipe_out, index], outputs=preview)
    output.change(show_final_preview, inputs=pipe_out, outputs=preview)
    gr.Markdown(
    """
    ## Advance usage:
    ### Browse restoration variants:
    1. You can explore other in-progress versions while dragging the "Preview Restoration" at the completion of SAKBIR Processing.
    2. You can get a more refined result by setting the "Start Free Rendering" for the output that you like.
    ### Creative restoration:
    1. You can also opt for the `Creative Restoration` option;
    2. Input the prompt in the textbox for "Restoration using prompts";
    3. Set the value of the slider for "Start Free Rendering" to a medium value for SAKBIR creation.
    """)

SAK.queue().launch()