Upload 4 files

app.py

@@ -1,154 +1,308 @@
-import gradio as gr
-import numpy as np
-import random
-
-import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-@spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-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(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+import numpy as np
+import random
+
+import spaces #[uncomment to use ZeroGPU]
+from diffusers import DiffusionPipeline
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+# model_repo_id = "/data/stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+#
+# if torch.cuda.is_available():
+#     torch_dtype = torch.float16
+# else:
+#     torch_dtype = torch.float32
+#
+# pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+# pipe = pipe.to(device)
+
+# ------------------ set up InterLCM restorer ------------------- #
+import os
+import cv2
+import argparse
+import glob
+import re
+import torch
+from torchvision.transforms.functional import normalize
+from basicsr.utils import imwrite, img2tensor, tensor2img
+from basicsr.utils.download_util import load_file_from_url
+from basicsr.utils.misc import gpu_is_available, get_device
+from facelib.utils.face_restoration_helper import FaceRestoreHelper
+from facelib.utils.misc import is_gray
+
+from basicsr.utils.registry import ARCH_REGISTRY
+
+# CILP
+import clip
+import torchvision.transforms as transforms
+
+from basicsr.utils.clip_util import VisionTransformer
+clip.model.VisionTransformer = VisionTransformer
+
+# LCM
+from diffusers import DiffusionPipeline, UNet2DConditionModel, ControlNetModel
+from basicsr.utils.lcm_utils import register_lcm_forward, register_lcmschedule_step
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.realesrgan_utils import RealESRGANer
+
+from scripts.wavelet_color_fix import wavelet_reconstruction, adaptive_instance_normalization
+
+visual_encoder_path = "weights/InterLCM/visual_encoder_3step.pth"
+spatial_encoder_path = "weights/InterLCM/spatial_encoder_3step.pth"
+visual_encoder_path_1step = "weights/InterLCM/visual_encoder_1step.pth"
+spatial_encoder_path_1step = "weights/InterLCM/spatial_encoder_1step.pth"
+sd_path = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+lcm_path = "SimianLuo/LCM_Dreamshaper_v7"
+detection_model = "retinaface_resnet50"
+
+# CLIPImageEncoder
+clip_model, clip_preprocess = clip.load('ViT-B/16', device=device)
+preprocess = transforms.Compose([transforms.Normalize(mean=[-1.0, -1.0, -1.0], std=[2.0, 2.0,
+                                                                                    2.0])] +  # Un-normalize from [-1.0, 1.0] (GAN output) to [0, 1].
+                                clip_preprocess.transforms[:2] +  # to match CLIP input scale assumptions
+                                clip_preprocess.transforms[4:])  # + skip convert PIL to tensor
+
+# Visual Encoder
+visual_encoder = ARCH_REGISTRY.get('VisualEncoder')(nf=64, emb_dim=197, ch_mult=[2, 4, 8], res_blocks=2,
+                                                    img_size=512).to(device)
+checkpoint_ve = torch.load(visual_encoder_path)['params_ema']
+visual_encoder.load_state_dict(checkpoint_ve)
+visual_encoder.eval()
+del checkpoint_ve
+
+# Spatial Encoder
+unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path=sd_path, subfolder="unet")
+spatial_encoder = ControlNetModel.from_unet(unet).to(device)
+checkpoint_c = torch.load(spatial_encoder_path)['params_ema']
+spatial_encoder.load_state_dict(checkpoint_c)
+spatial_encoder.eval()
+del unet
+
+# Visual Encoder 1-step
+visual_encoder_1step = ARCH_REGISTRY.get('VisualEncoder')(nf=64, emb_dim=197, ch_mult=[2, 4, 8], res_blocks=2,
+                                                    img_size=512).to(device)
+checkpoint_ve = torch.load(visual_encoder_path_1step)['params_ema']
+visual_encoder_1step.load_state_dict(checkpoint_ve)
+visual_encoder_1step.eval()
+del checkpoint_ve
+
+# Spatial Encoder
+unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path=sd_path, subfolder="unet")
+spatial_encoder_1step = ControlNetModel.from_unet(unet).to(device)
+checkpoint_c = torch.load(spatial_encoder_path_1step)['params_ema']
+spatial_encoder_1step.load_state_dict(checkpoint_c)
+spatial_encoder_1step.eval()
+del unet
+
+torch.cuda.empty_cache()
+
+# lcm
+lcm = DiffusionPipeline.from_pretrained(pretrained_model_name_or_path=lcm_path).to(device)
+
+# set enhancer with RealESRGAN
+def set_realesrgan():
+    half = True if torch.cuda.is_available() else False
+    model = RRDBNet(
+        num_in_ch=3,
+        num_out_ch=3,
+        num_feat=64,
+        num_block=23,
+        num_grow_ch=32,
+        scale=2,
+    )
+    upsampler = RealESRGANer(
+        scale=2,
+        model_path="weights/realesrgan/RealESRGAN_x2plus.pth",
+        model=model,
+        tile=400,
+        tile_pad=40,
+        pre_pad=0,
+        half=half,
+        device=device
+    )
+    return upsampler
+
+upsampler = set_realesrgan()
+
+upscale = 2
+face_helper = FaceRestoreHelper(
+    upscale_factor=upscale,
+    face_size=512,
+    crop_ratio=(1, 1),
+    det_model=detection_model,
+    save_ext='png',
+    use_parse=True,
+    device=device)
+
+# ------------------ set up InterLCM restorer ------------------- #
+@spaces.GPU
+def inference(input_img, interlcm_step, face_align, background_enhance, face_upsample):
+    # try:
+        only_center_face = False
+        draw_box = False
+
+        interlcm_step = int(interlcm_step)
+        assert interlcm_step in (1, 3)
+        if interlcm_step == 1:
+            register_lcm_forward(lcm, spatial_encoder_1step)
+        elif interlcm_step == 3:
+            register_lcm_forward(lcm, spatial_encoder)
+        register_lcmschedule_step(lcm.scheduler)
+
+        face_align = face_align if face_align is not None else True
+        has_aligned = not face_align
+
+        background_enhance = background_enhance if background_enhance is not None else True
+        bg_upsampler = upsampler if background_enhance else None
+
+        face_upsampler = upsampler if face_upsample else None
+
+        img = cv2.imread(str(input_img), cv2.IMREAD_COLOR)
+        print('\timage size:', img.shape)
+
+        face_helper.clean_all()
+        if has_aligned:
+            # the input faces are already cropped and aligned
+            img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
+            face_helper.is_gray = is_gray(img, threshold=10)
+            if face_helper.is_gray:
+                print('Grayscale input: True')
+            face_helper.cropped_faces = [img]
+        else:
+            face_helper.read_image(img)
+            # get face landmarks for each face
+            num_det_faces = face_helper.get_face_landmarks_5(
+                only_center_face=only_center_face, resize=640, eye_dist_threshold=5, device=device)
+            print(f'\tdetect {num_det_faces} faces')
+            # align and warp each face
+            face_helper.align_warp_face()
+
+        # face restoration for each cropped face
+        for idx, cropped_face in enumerate(face_helper.cropped_faces):
+            # prepare data
+            cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
+            normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+            cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
+
+            try:
+                with torch.no_grad():
+                    input = preprocess(cropped_face_t)
+                    img_emb = clip_model.encode_image(input)
+                    img_emb = img_emb.to(torch.float)
+
+                    if interlcm_step == 1:
+                        visual_feat = visual_encoder_1step(img_emb)
+                    elif interlcm_step == 3:
+                        visual_feat = visual_encoder(img_emb)
+
+                    latent_code = lcm.vae.encode(cropped_face_t)['latent_dist'].mean
+                    latent_code = latent_code * 0.18215
+                    output = lcm.forward(height=512, width=512, num_inference_steps=interlcm_step + 1,
+                                         guidance_scale=8.0, latents=latent_code,
+                                         prompt_embeds=visual_feat, output_type="pil", lcm_origin_steps=50,
+                                         lq_input=cropped_face_t).images
+                    output = wavelet_reconstruction(output, cropped_face_t)
+                    restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
+
+                    del output
+                    torch.cuda.empty_cache()
+            except Exception as error:
+                print(f'\tFailed inference for CodeFormer: {error}')
+                restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
+
+            restored_face = restored_face.astype('uint8')
+            face_helper.add_restored_face(restored_face, cropped_face)
+
+        # paste_back
+        if not has_aligned:
+            # upsample the background
+            if bg_upsampler is not None:
+                # Now only support RealESRGAN for upsampling background
+                bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
+            else:
+                bg_img = None
+            face_helper.get_inverse_affine(None)
+            # paste each restored face to the input image
+            if face_upsample and face_upsampler is not None:
+                restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=draw_box,
+                                                                      face_upsampler=face_upsampler)
+            else:
+                restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=draw_box)
+        else:
+            restored_img = restored_face
+
+        # save restored img
+        save_path = f'output/out.png'
+        imwrite(restored_img, save_path)
+
+        restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
+
+        return restored_img
+    # except Exception as error:
+    #     print('Global exception', error)
+    #     return None
+
+
+title = "InterLCM: Low-Quality Images as Intermediate States of Latent Consistency Models for Effective Blind Face Restoration"
+
+description = r"""<center><img src='https://raw.githubusercontent.com/sen-mao/InterLCM/refs/heads/master/assets/interlcm_logo.jpg' alt='InterLCM logo' width="120"></center>
+<br>
+<b>Official Gradio demo</b> for <a href='https://github.com/sen-mao/InterLCM' target='_blank'><b>Low-Quality Images as Intermediate States of Latent Consistency Models for Effective Blind Face Restoration (ICLR 2025)</b></a><br>
+🔥 InterLCM is a robust blind face restoration algorithm.<br>
+⭐ If InterLCM is helpful to your images or projects, please help star this repo. Thanks! 🤗 <br>
+"""
+
+article = r"""
+If InterLCM is helpful, please help to ⭐ the <a href='https://github.com/sen-mao/InterLCM' target='_blank'>Github Repo</a>. Thanks! 
+[![GitHub Stars](https://img.shields.io/github/stars/sen-mao/InterLCM?style=social)](https://github.com/sen-mao/InterLCM)
+
+---
+
+📝 **Citation**
+If our work is useful for your research, please consider citing:
+```bibtex
+@inproceedings{li2025interlcm,
+  title={InterLCM: Low-Quality Images as Intermediate States of Latent Consistency Models for Effective Blind Face Restoration},
+  author={Li, Senmao and Wang, Kai and van de Weijer, Joost and Khan, Fahad Shahbaz and Guo, Chun-Le and Yang, Shiqi and Wang, Yaxing and Yang, Jian and Cheng, Ming-Ming},
+  booktitle={ICLR},
+  year={2025}
+}
+```
+
+📧 **Contact**
+If you have any questions, please feel free to reach me out at <b>senmaonk@gmail.com</b>.
+
+<center><img src='https://visitor-badge.laobi.icu/badge?page_id=sen-mao/InterLCM&ltext=Visitors' alt='visitors'></center>
+"""
+
+
+demo = gr.Interface(
+    inference, [
+        gr.Image(type="filepath", label="Input"),
+        gr.Radio(choices=["1", "3"], value="3", label="Select InterLCM step (InterLCM enables 1-step⚡ BFR under non-extreme degradation conditions)"),
+        gr.Checkbox(value=True, label="Pre_Face_Align"),
+        gr.Checkbox(value=True, label="Background_Enhance"),
+        gr.Checkbox(value=True, label="Face_Upsample"),
+    ], [
+        gr.Image(type="numpy", label="Output")
+    ],
+    title=title,
+    description=description,
+    article=article,
+    examples=[
+        ['inputs/cropped_faces/0631.png', "3", False, False, False],
+        ['inputs/cropped_faces/Nora_Bendijo_0001_00.png', "3", False, False, False],
+        ['inputs/whole_imgs/03.jpg', "1", True, True, True],
+        ['inputs/whole_imgs/04.jpg', "3", True, True, True],
+        ['inputs/whole_imgs/05.jpg', "3", True, True, True]
+    ],
+    concurrency_limit=2,
+    allow_flagging="never",
+)
+
+
+if __name__ == "__main__":
+    # DEBUG = os.getenv('DEBUG') == '1'
+    # demo.launch(server_name="0.0.0.0", server_port=7861, max_threads=10, share=False)
+    demo.launch()

inference_InterLCM.py

@@ -0,0 +1,346 @@
+import os
+import cv2
+import argparse
+import glob
+import re
+import torch
+from torchvision.transforms.functional import normalize
+from basicsr.utils import imwrite, img2tensor, tensor2img
+from basicsr.utils.download_util import load_file_from_url
+from basicsr.utils.misc import gpu_is_available, get_device
+from facelib.utils.face_restoration_helper import FaceRestoreHelper
+from facelib.utils.misc import is_gray
+
+from basicsr.utils.registry import ARCH_REGISTRY
+
+# CILP
+import clip
+import torchvision.transforms as transforms
+
+from basicsr.utils.clip_util import VisionTransformer
+clip.model.VisionTransformer = VisionTransformer
+
+# LCM
+from diffusers import DiffusionPipeline, UNet2DConditionModel, ControlNetModel
+from basicsr.utils.lcm_utils import register_lcm_forward, register_lcmschedule_step
+
+from scripts.wavelet_color_fix import wavelet_reconstruction, adaptive_instance_normalization
+
+def set_realesrgan(args):
+    from basicsr.archs.rrdbnet_arch import RRDBNet
+    from basicsr.utils.realesrgan_utils import RealESRGANer
+
+    use_half = False
+    if torch.cuda.is_available(): # set False in CPU/MPS mode
+        no_half_gpu_list = ['1650', '1660'] # set False for GPUs that don't support f16
+        if not True in [gpu in torch.cuda.get_device_name(0) for gpu in no_half_gpu_list]:
+            use_half = True
+
+    model = RRDBNet(
+        num_in_ch=3,
+        num_out_ch=3,
+        num_feat=64,
+        num_block=23,
+        num_grow_ch=32,
+        scale=2,
+    )
+    upsampler = RealESRGANer(
+        scale=2,
+        model_path="https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth",
+        model=model,
+        tile=args.bg_tile,
+        tile_pad=40,
+        pre_pad=0,
+        half=use_half
+    )
+
+    if not gpu_is_available():  # CPU
+        import warnings
+        warnings.warn('Running on CPU now! Make sure your PyTorch version matches your CUDA.'
+                        'The unoptimized RealESRGAN is slow on CPU. '
+                        'If you want to disable it, please remove `--bg_upsampler` and `--face_upsample` in command.',
+                        category=RuntimeWarning)
+    return upsampler
+
+@torch.no_grad()
+def main():
+    device = get_device()
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('-i', '--input_path', type=str, default='./inputs/whole_imgs', 
+            help='Input image, video or folder. Default: inputs/whole_imgs')
+    parser.add_argument('-o', '--output_path', type=str, default="results",
+            help='Output folder. Default: results/<input_name>')
+    parser.add_argument('-s', '--upscale', type=int, default=2, 
+            help='The final upsampling scale of the image. Default: 2')
+    parser.add_argument('--has_aligned', action='store_true', help='Input are cropped and aligned faces. Default: False')
+    parser.add_argument('--only_center_face', action='store_true', help='Only restore the center face. Default: False')
+    parser.add_argument('--draw_box', action='store_true', help='Draw the bounding box for the detected faces. Default: False')
+    # large det_model: 'YOLOv5l', 'retinaface_resnet50'
+    # small det_model: 'YOLOv5n', 'retinaface_mobile0.25'
+    parser.add_argument('--detection_model', type=str, default='retinaface_resnet50', 
+            help='Face detector. Optional: retinaface_resnet50, retinaface_mobile0.25, YOLOv5l, YOLOv5n, dlib. \
+                Default: retinaface_resnet50')
+    parser.add_argument('--bg_upsampler', type=str, default='None', help='Background upsampler. Optional: realesrgan')
+    parser.add_argument('--face_upsample', action='store_true', help='Face upsampler after enhancement. Default: False')
+    parser.add_argument('--bg_tile', type=int, default=400, help='Tile size for background sampler. Default: 400')
+    parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces. Default: None')
+    parser.add_argument('--save_video_fps', type=float, default=None, help='Frame rate for saving video. Default: None')
+    # LCM
+    parser.add_argument('--num_inference_steps', type=int, default=4, help='T for lcm')
+    parser.add_argument('--visual_encoder_path', type=str,
+                        default='weights/InterLCM/visual_encoder_3step.pth',
+                        help='visual_encoder checkpoint')
+    parser.add_argument('--spatial_encoder_path', type=str,
+                        default='weights/InterLCM/spatial_encoder_3step.pth',
+                        help='spatial_encoder checkpoint')
+    parser.add_argument('--sd_path', type=str,
+                        default='/data/runwayml/stable-diffusion-v1-5',
+                        help='sd pre-trined model')
+    parser.add_argument('--lcm_path', type=str,
+                        default='/data/SimianLuo/LCM_Dreamshaper_v7',
+                        help='lcm pre-trined model')
+
+    parser.add_argument(
+        "--colorfix_type",
+        type=str,
+        default="wavelet",
+        help="Color fix type to adjust the color of reconstructed HR result according to LR input: "
+             "adain; wavelet (used in paper); nofix",
+    )
+
+    args = parser.parse_args()
+    print(args)
+
+    interlcm_step = int(re.findall(r'\d+', args.visual_encoder_path)[0])
+    args.output_path = args.output_path.replace(args.output_path.split('/')[0],
+                                                f"{args.output_path.split('/')[0]}[{args.colorfix_type}]/interlcm_{interlcm_step}step")
+
+    assert args.num_inference_steps - 1 == interlcm_step
+
+    # ------------------------ input & output ------------------------
+    input_video = False
+    if args.input_path.endswith(('jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG')): # input single img path
+        input_img_list = [args.input_path]
+        result_root = f'results/test_img'
+    elif args.input_path.endswith(('mp4', 'mov', 'avi', 'MP4', 'MOV', 'AVI')): # input video path
+        from basicsr.utils.video_util import VideoReader, VideoWriter
+        input_img_list = []
+        vidreader = VideoReader(args.input_path)
+        image = vidreader.get_frame()
+        while image is not None:
+            input_img_list.append(image)
+            image = vidreader.get_frame()
+        audio = vidreader.get_audio()
+        fps = vidreader.get_fps() if args.save_video_fps is None else args.save_video_fps   
+        video_name = os.path.basename(args.input_path)[:-4]
+        result_root = f'results/{video_name}'
+        input_video = True
+        vidreader.close()
+    else: # input img folder
+        if args.input_path.endswith('/'):  # solve when path ends with /
+            args.input_path = args.input_path[:-1]
+        # scan all the jpg and png images
+        input_img_list = sorted(glob.glob(os.path.join(args.input_path, '*.[jpJP][pnPN]*[gG]')))
+        result_root = f'results/{os.path.basename(args.input_path)}'
+
+    if not args.output_path is None: # set output path
+        result_root = args.output_path
+
+    test_img_num = len(input_img_list)
+    if test_img_num == 0:
+        raise FileNotFoundError('No input image/video is found...\n' 
+            '\tNote that --input_path for video should end with .mp4|.mov|.avi')
+
+    # ------------------ set up background upsampler ------------------
+    if args.bg_upsampler == 'realesrgan':
+        bg_upsampler = set_realesrgan(args)
+    else:
+        bg_upsampler = None
+
+    # ------------------ set up face upsampler ------------------
+    if args.face_upsample:
+        if bg_upsampler is not None:
+            face_upsampler = bg_upsampler
+        else:
+            face_upsampler = set_realesrgan(args)
+    else:
+        face_upsampler = None
+
+    # ------------------ set up InterLCM restorer -------------------
+
+    # CLIPImageEncoder
+    clip_model, clip_preprocess = clip.load('ViT-B/16', device=device)
+    preprocess = transforms.Compose([transforms.Normalize(mean=[-1.0, -1.0, -1.0], std=[2.0, 2.0, 2.0])] +  # Un-normalize from [-1.0, 1.0] (GAN output) to [0, 1].
+                                    clip_preprocess.transforms[:2] +  # to match CLIP input scale assumptions
+                                    clip_preprocess.transforms[4:])  # + skip convert PIL to tensor
+
+    # Visual Encoder
+    visual_encoder = ARCH_REGISTRY.get('VisualEncoder')(nf=64, emb_dim=197, ch_mult=[2,4,8], res_blocks=2, img_size=512).to(device)
+    checkpoint_ve = torch.load(args.visual_encoder_path)['params_ema']
+    visual_encoder.load_state_dict(checkpoint_ve)
+    visual_encoder.eval()
+
+    # Spatial Encoder
+    unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path=args.sd_path, subfolder="unet")
+    spatial_encoder = ControlNetModel.from_unet(unet).to(device)
+    checkpoint_c = torch.load(args.spatial_encoder_path)['params_ema']
+    spatial_encoder.load_state_dict(checkpoint_c)
+    spatial_encoder.eval()
+
+    # lcm
+    lcm = DiffusionPipeline.from_pretrained(pretrained_model_name_or_path=args.lcm_path).to(device)
+
+    register_lcm_forward(lcm, spatial_encoder)
+    register_lcmschedule_step(lcm.scheduler)
+
+
+    # ------------------ set up FaceRestoreHelper -------------------
+    # large det_model: 'YOLOv5l', 'retinaface_resnet50'
+    # small det_model: 'YOLOv5n', 'retinaface_mobile0.25'
+    if not args.has_aligned: 
+        print(f'Face detection model: {args.detection_model}')
+    if bg_upsampler is not None: 
+        print(f'Background upsampling: True, Face upsampling: {args.face_upsample}')
+    else:
+        print(f'Background upsampling: False, Face upsampling: {args.face_upsample}')
+
+    face_helper = FaceRestoreHelper(
+        args.upscale,
+        face_size=512,
+        crop_ratio=(1, 1),
+        det_model = args.detection_model,
+        save_ext='png',
+        use_parse=True,
+        device=device)
+
+    # -------------------- start to processing ---------------------
+    for i, img_path in enumerate(input_img_list):
+        # clean all the intermediate results to process the next image
+        face_helper.clean_all()
+        
+        if isinstance(img_path, str):
+            img_name = os.path.basename(img_path)
+            basename, ext = os.path.splitext(img_name)
+            print(f'[{i+1}/{test_img_num}] Processing: {img_name}')
+            img = cv2.imread(img_path, cv2.IMREAD_COLOR)
+        else: # for video processing
+            basename = str(i).zfill(6)
+            img_name = f'{video_name}_{basename}' if input_video else basename
+            print(f'[{i+1}/{test_img_num}] Processing: {img_name}')
+            img = img_path
+
+        if args.has_aligned: 
+            # the input faces are already cropped and aligned
+            img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
+            face_helper.is_gray = is_gray(img, threshold=10)
+            if face_helper.is_gray:
+                print('Grayscale input: True')
+            face_helper.cropped_faces = [img]
+        else:
+            face_helper.read_image(img)
+            # get face landmarks for each face
+            num_det_faces = face_helper.get_face_landmarks_5(
+                only_center_face=args.only_center_face, resize=640, eye_dist_threshold=5)
+            print(f'\tdetect {num_det_faces} faces')
+            # align and warp each face
+            face_helper.align_warp_face()
+
+        # face restoration for each cropped face
+        for idx, cropped_face in enumerate(face_helper.cropped_faces):
+            # prepare data
+            cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
+            normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+            cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
+
+            try:
+                with torch.no_grad():
+                    input = preprocess(cropped_face_t)
+                    img_emb = clip_model.encode_image(input)
+                    img_emb = img_emb.to(torch.float)
+
+                    visual_feat = visual_encoder(img_emb)
+
+                    latent_code = lcm.vae.encode(cropped_face_t)['latent_dist'].mean
+                    latent_code = latent_code * 0.18215
+                    output = lcm.forward(height=512, width=512, num_inference_steps=args.num_inference_steps, guidance_scale=8.0, latents=latent_code,
+                                         prompt_embeds=visual_feat, output_type="pil", lcm_origin_steps=50, lq_input=cropped_face_t).images
+
+                    # colorfix from StableSR
+                    if args.colorfix_type == 'adain':
+                        output = adaptive_instance_normalization(output, cropped_face_t)
+                    elif args.colorfix_type == 'wavelet':
+                        output = wavelet_reconstruction(output, cropped_face_t)
+
+                    restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
+                del output
+                torch.cuda.empty_cache()
+            except Exception as error:
+                print(f'\tFailed inference for CodeFormer: {error}')
+                restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
+
+            restored_face = restored_face.astype('uint8')
+            face_helper.add_restored_face(restored_face, cropped_face)
+
+        # paste_back
+        if not args.has_aligned:
+            # upsample the background
+            if bg_upsampler is not None:
+                # Now only support RealESRGAN for upsampling background
+                bg_img = bg_upsampler.enhance(img, outscale=args.upscale)[0]
+            else:
+                bg_img = None
+            face_helper.get_inverse_affine(None)
+            # paste each restored face to the input image
+            if args.face_upsample and face_upsampler is not None: 
+                restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=args.draw_box, face_upsampler=face_upsampler)
+            else:
+                restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=args.draw_box)
+
+        # save faces
+        for idx, (cropped_face, restored_face) in enumerate(zip(face_helper.cropped_faces, face_helper.restored_faces)):
+            # save cropped face
+            if not args.has_aligned: 
+                save_crop_path = os.path.join(result_root, 'cropped_faces', f'{basename}_{idx:02d}.png')
+                imwrite(cropped_face, save_crop_path)
+            # save restored face
+            if args.has_aligned:
+                save_face_name = f'{basename}.png'
+            else:
+                save_face_name = f'{basename}_{idx:02d}.png'
+            if args.suffix is not None:
+                save_face_name = f'{save_face_name[:-4]}_{args.suffix}.png'
+            save_restore_path = os.path.join(result_root, 'restored_faces', save_face_name)
+            imwrite(restored_face, save_restore_path)
+
+        # save restored img
+        if not args.has_aligned and restored_img is not None:
+            if args.suffix is not None:
+                basename = f'{basename}_{args.suffix}'
+            save_restore_path = os.path.join(result_root, 'final_results', f'{basename}.png')
+            imwrite(restored_img, save_restore_path)
+
+    # save enhanced video
+    if input_video:
+        print('Video Saving...')
+        # load images
+        video_frames = []
+        img_list = sorted(glob.glob(os.path.join(result_root, 'final_results', '*.[jp][pn]g')))
+        for img_path in img_list:
+            img = cv2.imread(img_path)
+            video_frames.append(img)
+        # write images to video
+        height, width = video_frames[0].shape[:2]
+        if args.suffix is not None:
+            video_name = f'{video_name}_{args.suffix}.png'
+        save_restore_path = os.path.join(result_root, f'{video_name}.mp4')
+        vidwriter = VideoWriter(save_restore_path, height, width, fps, audio)
+         
+        for f in video_frames:
+            vidwriter.write_frame(f)
+        vidwriter.close()
+
+    print(f'\nAll results are saved in {result_root}')
+
+if __name__ == '__main__':
+    main()

sdxlturbo.py

@@ -0,0 +1,154 @@
+import gradio as gr
+import numpy as np
+import random
+
+# import spaces #[uncomment to use ZeroGPU]
+from diffusers import DiffusionPipeline
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_repo_id = "/data/stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+
+if torch.cuda.is_available():
+    torch_dtype = torch.float16
+else:
+    torch_dtype = torch.float32
+
+pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+pipe = pipe.to(device)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 512
+
+
+# @spaces.GPU #[uncomment to use ZeroGPU]
+def infer(
+    prompt,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    width,
+    height,
+    guidance_scale,
+    num_inference_steps,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    generator = torch.Generator().manual_seed(seed)
+
+    image = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        generator=generator,
+    ).images[0]
+
+    return image, seed
+
+
+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(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(" # InterLCM: Low-Quality Images as Intermediate States of Latent Consistency Models for Effective Blind Face Restoration")
+
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+
+            run_button = gr.Button("Run", scale=0, variant="primary")
+
+        result = gr.Image(label="Result", show_label=False)
+
+        with gr.Accordion("Advanced Settings", open=False):
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+                visible=False,
+            )
+
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+            with gr.Row():
+                width = gr.Slider(
+                    label="Width",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=1024,  # Replace with defaults that work for your model
+                )
+
+                height = gr.Slider(
+                    label="Height",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=1024,  # Replace with defaults that work for your model
+                )
+
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=0.0,  # Replace with defaults that work for your model
+                )
+
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=2,  # Replace with defaults that work for your model
+                )
+
+        gr.Examples(examples=examples, inputs=[prompt])
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[
+            prompt,
+            negative_prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps,
+        ],
+        outputs=[result, seed],
+    )
+
+if __name__ == "__main__":
+    demo.launch()

sepia.py

@@ -0,0 +1,18 @@
+import numpy as np
+import gradio as gr
+
+
+def sepia(input_img):
+    sepia_filter = np.array([
+        [0.393, 0.769, 0.189],
+        [0.349, 0.686, 0.168],
+        [0.272, 0.534, 0.131]
+    ])
+    sepia_img = input_img.dot(sepia_filter.T)
+    sepia_img /= sepia_img.max()
+    return sepia_img
+
+
+gr.ChatInterface(sepia, analytics_enabled=False)
+demo = gr.Interface(sepia, gr.Image(), "image")
+demo.launch(share=True)