Upload gradio_seesr.py

gradio_seesr.py

@@ -0,0 +1,207 @@
+import gradio as gr
+import os
+import sys
+from typing import List
+# sys.path.append(os.getcwd())
+
+import numpy as np
+from PIL import Image
+
+import torch
+import torch.utils.checkpoint
+from pytorch_lightning import seed_everything
+from diffusers import AutoencoderKL, DDPMScheduler
+from diffusers.utils import check_min_version
+from diffusers.utils.import_utils import is_xformers_available
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor
+
+from pipelines.pipeline_seesr import StableDiffusionControlNetPipeline
+
+from utils.wavelet_color_fix import wavelet_color_fix, adain_color_fix
+
+from ram.models.ram_lora import ram
+from ram import inference_ram as inference
+from torchvision import transforms
+from models.controlnet import ControlNetModel
+from models.unet_2d_condition import UNet2DConditionModel
+
+tensor_transforms = transforms.Compose([
+                transforms.ToTensor(),
+            ])
+
+ram_transforms = transforms.Compose([
+            transforms.Resize((384, 384)),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+
+# Load scheduler, tokenizer and models.
+pretrained_model_path = 'preset/models/stable-diffusion-2-1-base'
+seesr_model_path = 'preset/models/seesr'
+
+scheduler = DDPMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
+text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder")
+tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
+vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae")
+feature_extractor = CLIPImageProcessor.from_pretrained(f"{pretrained_model_path}/feature_extractor")
+unet = UNet2DConditionModel.from_pretrained(seesr_model_path, subfolder="unet")
+controlnet = ControlNetModel.from_pretrained(seesr_model_path, subfolder="controlnet")
+
+# Freeze vae and text_encoder
+vae.requires_grad_(False)
+text_encoder.requires_grad_(False)
+unet.requires_grad_(False)
+controlnet.requires_grad_(False)
+
+if is_xformers_available():
+    unet.enable_xformers_memory_efficient_attention()
+    controlnet.enable_xformers_memory_efficient_attention()
+else:
+    raise ValueError("xformers is not available. Make sure it is installed correctly")
+
+# Get the validation pipeline
+validation_pipeline = StableDiffusionControlNetPipeline(
+    vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, feature_extractor=feature_extractor,
+    unet=unet, controlnet=controlnet, scheduler=scheduler, safety_checker=None, requires_safety_checker=False,
+)
+
+validation_pipeline._init_tiled_vae(encoder_tile_size=1024,
+                                    decoder_tile_size=224)
+weight_dtype = torch.float16
+device = "cuda"
+
+# Move text_encode and vae to gpu and cast to weight_dtype
+text_encoder.to(device, dtype=weight_dtype)
+vae.to(device, dtype=weight_dtype)
+unet.to(device, dtype=weight_dtype)
+controlnet.to(device, dtype=weight_dtype)
+
+
+tag_model = ram(pretrained='preset/models/ram_swin_large_14m.pth',
+                pretrained_condition='preset/models/DAPE.pth',
+                image_size=384,
+                vit='swin_l')
+tag_model.eval()
+tag_model.to(device, dtype=weight_dtype)
+
+@torch.no_grad()
+def process(
+    input_image: Image.Image,
+    user_prompt: str,
+    positive_prompt: str,
+    negative_prompt: str,
+    num_inference_steps: int,
+    scale_factor: int,
+    cfg_scale: float,
+    seed: int,
+    latent_tiled_size: int,
+    latent_tiled_overlap: int,
+    sample_times: int
+    ) -> List[np.ndarray]:
+    process_size = 512
+    resize_preproc = transforms.Compose([
+        transforms.Resize(process_size, interpolation=transforms.InterpolationMode.BILINEAR),
+    ])
+
+    # with torch.no_grad():
+    seed_everything(seed)
+    generator = torch.Generator(device=device)
+
+    validation_prompt = ""
+    lq = tensor_transforms(input_image).unsqueeze(0).to(device).half()
+    lq = ram_transforms(lq)
+    res = inference(lq, tag_model)
+    ram_encoder_hidden_states = tag_model.generate_image_embeds(lq)
+    validation_prompt = f"{res[0]}, {positive_prompt},"
+    validation_prompt = validation_prompt if user_prompt=='' else f"{user_prompt}, {validation_prompt}"
+
+    ori_width, ori_height = input_image.size
+    resize_flag = False
+
+    rscale = scale_factor
+    input_image = input_image.resize((int(input_image.size[0] * rscale), int(input_image.size[1] * rscale)))
+
+    if min(input_image.size) < process_size:
+        input_image = resize_preproc(input_image)
+
+    input_image = input_image.resize((input_image.size[0] // 8 * 8, input_image.size[1] // 8 * 8))
+    width, height = input_image.size
+    resize_flag = True  #
+
+    images = []
+    for _ in range(sample_times):
+        try:
+            with torch.autocast("cuda"):
+                image = validation_pipeline(
+                    validation_prompt, input_image, negative_prompt=negative_prompt,
+                    num_inference_steps=num_inference_steps, generator=generator,
+                    height=height, width=width,
+                    guidance_scale=cfg_scale,  conditioning_scale=1,
+                    start_point='lr', start_steps=999,ram_encoder_hidden_states=ram_encoder_hidden_states,
+                    latent_tiled_size=latent_tiled_size, latent_tiled_overlap=latent_tiled_overlap
+                ).images[0]
+
+            if True:  # alpha<1.0:
+                image = wavelet_color_fix(image, input_image)
+
+            if resize_flag:
+                image = image.resize((ori_width * rscale, ori_height * rscale))
+        except Exception as e:
+            print(e)
+            image = Image.new(mode="RGB", size=(512, 512))
+        images.append(np.array(image))
+    return images
+
+
+#
+MARKDOWN = \
+"""
+## SeeSR: Towards Semantics-Aware Real-World Image Super-Resolution
+
+[GitHub](https://github.com/cswry/SeeSR) | [Paper](https://arxiv.org/abs/2311.16518)
+
+If SeeSR is helpful for you, please help star the GitHub Repo. Thanks!
+"""
+
+block = gr.Blocks().queue()
+with block:
+    with gr.Row():
+        gr.Markdown(MARKDOWN)
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(source="upload", type="pil")
+            run_button = gr.Button(label="Run")
+            with gr.Accordion("Options", open=True):
+                user_prompt = gr.Textbox(label="User Prompt", value="")
+                positive_prompt = gr.Textbox(label="Positive Prompt", value="clean, high-resolution, 8k, best quality, masterpiece")
+                negative_prompt = gr.Textbox(
+                    label="Negative Prompt",
+                    value="dotted, noise, blur, lowres, oversmooth, longbody, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
+                )
+                cfg_scale = gr.Slider(label="Classifier Free Guidance Scale (Set a value larger than 1 to enable it!)", minimum=0.1, maximum=10.0, value=5.5, step=0.1)
+                num_inference_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=50, step=1)
+                seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=231)
+                sample_times = gr.Slider(label="Sample Times", minimum=1, maximum=10, step=1, value=1)
+                latent_tiled_size = gr.Slider(label="Diffusion Tile Size", minimum=128, maximum=480, value=320, step=1)
+                latent_tiled_overlap = gr.Slider(label="Diffusion Tile Overlap", minimum=4, maximum=16, value=4, step=1)
+                scale_factor = gr.Number(label="SR Scale", value=4)
+        with gr.Column():
+            result_gallery = gr.Gallery(label="Output", show_label=False, elem_id="gallery").style(grid=2, height="auto")
+
+    inputs = [
+        input_image,
+        user_prompt,
+        positive_prompt,
+        negative_prompt,
+        num_inference_steps,
+        scale_factor,
+        cfg_scale,
+        seed,
+        latent_tiled_size,
+        latent_tiled_overlap,
+        sample_times,
+    ]
+    run_button.click(fn=process, inputs=inputs, outputs=[result_gallery])
+
+block.launch()
+