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()