ProductScope/Relight

1import spaces2import math3import gradio as gr4import numpy as np5import torch6import safetensors.torch as sf7import db_examples8 9from PIL import Image10from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline11from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler, EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler12from diffusers.models.attention_processor import AttnProcessor2_013from transformers import CLIPTextModel, CLIPTokenizer14from briarmbg import BriaRMBG15from enum import Enum16# from torch.hub import download_url_to_file17 18 19# 'stablediffusionapi/realistic-vision-v51'20# 'runwayml/stable-diffusion-v1-5'21sd15_name = 'stablediffusionapi/realistic-vision-v51'22tokenizer = CLIPTokenizer.from_pretrained(sd15_name, subfolder="tokenizer")23text_encoder = CLIPTextModel.from_pretrained(sd15_name, subfolder="text_encoder")24vae = AutoencoderKL.from_pretrained(sd15_name, subfolder="vae")25unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet")26rmbg = BriaRMBG.from_pretrained("briaai/RMBG-1.4")27 28# Change UNet29 30with torch.no_grad():31    new_conv_in = torch.nn.Conv2d(8, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding)32    new_conv_in.weight.zero_()33    new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)34    new_conv_in.bias = unet.conv_in.bias35    unet.conv_in = new_conv_in36 37unet_original_forward = unet.forward38 39 40def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs):41    c_concat = kwargs['cross_attention_kwargs']['concat_conds'].to(sample)42    c_concat = torch.cat([c_concat] * (sample.shape[0] // c_concat.shape[0]), dim=0)43    new_sample = torch.cat([sample, c_concat], dim=1)44    kwargs['cross_attention_kwargs'] = {}45    return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs)46 47 48unet.forward = hooked_unet_forward49 50# Load51 52model_path = './models/iclight_sd15_fc.safetensors'53# download_url_to_file(url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fc.safetensors', dst=model_path)54sd_offset = sf.load_file(model_path)55sd_origin = unet.state_dict()56keys = sd_origin.keys()57sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}58unet.load_state_dict(sd_merged, strict=True)59del sd_offset, sd_origin, sd_merged, keys60 61# Device62 63device = torch.device('cuda')64text_encoder = text_encoder.to(device=device, dtype=torch.float16)65vae = vae.to(device=device, dtype=torch.bfloat16)66unet = unet.to(device=device, dtype=torch.float16)67rmbg = rmbg.to(device=device, dtype=torch.float32)68 69# SDP70 71unet.set_attn_processor(AttnProcessor2_0())72vae.set_attn_processor(AttnProcessor2_0())73 74# Samplers75 76ddim_scheduler = DDIMScheduler(77    num_train_timesteps=1000,78    beta_start=0.00085,79    beta_end=0.012,80    beta_schedule="scaled_linear",81    clip_sample=False,82    set_alpha_to_one=False,83    steps_offset=1,84)85 86euler_a_scheduler = EulerAncestralDiscreteScheduler(87    num_train_timesteps=1000,88    beta_start=0.00085,89    beta_end=0.012,90    steps_offset=191)92 93dpmpp_2m_sde_karras_scheduler = DPMSolverMultistepScheduler(94    num_train_timesteps=1000,95    beta_start=0.00085,96    beta_end=0.012,97    algorithm_type="sde-dpmsolver++",98    use_karras_sigmas=True,99    steps_offset=1100)101 102# Pipelines103 104t2i_pipe = StableDiffusionPipeline(105    vae=vae,106    text_encoder=text_encoder,107    tokenizer=tokenizer,108    unet=unet,109    scheduler=dpmpp_2m_sde_karras_scheduler,110    safety_checker=None,111    requires_safety_checker=False,112    feature_extractor=None,113    image_encoder=None114)115 116i2i_pipe = StableDiffusionImg2ImgPipeline(117    vae=vae,118    text_encoder=text_encoder,119    tokenizer=tokenizer,120    unet=unet,121    scheduler=dpmpp_2m_sde_karras_scheduler,122    safety_checker=None,123    requires_safety_checker=False,124    feature_extractor=None,125    image_encoder=None126)127 128 129@torch.inference_mode()130def encode_prompt_inner(txt: str):131    max_length = tokenizer.model_max_length132    chunk_length = tokenizer.model_max_length - 2133    id_start = tokenizer.bos_token_id134    id_end = tokenizer.eos_token_id135    id_pad = id_end136 137    def pad(x, p, i):138        return x[:i] if len(x) >= i else x + [p] * (i - len(x))139 140    tokens = tokenizer(txt, truncation=False, add_special_tokens=False)["input_ids"]141    chunks = [[id_start] + tokens[i: i + chunk_length] + [id_end] for i in range(0, len(tokens), chunk_length)]142    chunks = [pad(ck, id_pad, max_length) for ck in chunks]143 144    token_ids = torch.tensor(chunks).to(device=device, dtype=torch.int64)145    conds = text_encoder(token_ids).last_hidden_state146 147    return conds148 149 150@torch.inference_mode()151def encode_prompt_pair(positive_prompt, negative_prompt):152    c = encode_prompt_inner(positive_prompt)153    uc = encode_prompt_inner(negative_prompt)154 155    c_len = float(len(c))156    uc_len = float(len(uc))157    max_count = max(c_len, uc_len)158    c_repeat = int(math.ceil(max_count / c_len))159    uc_repeat = int(math.ceil(max_count / uc_len))160    max_chunk = max(len(c), len(uc))161 162    c = torch.cat([c] * c_repeat, dim=0)[:max_chunk]163    uc = torch.cat([uc] * uc_repeat, dim=0)[:max_chunk]164 165    c = torch.cat([p[None, ...] for p in c], dim=1)166    uc = torch.cat([p[None, ...] for p in uc], dim=1)167 168    return c, uc169 170 171@torch.inference_mode()172def pytorch2numpy(imgs, quant=True):173    results = []174    for x in imgs:175        y = x.movedim(0, -1)176 177        if quant:178            y = y * 127.5 + 127.5179            y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)180        else:181            y = y * 0.5 + 0.5182            y = y.detach().float().cpu().numpy().clip(0, 1).astype(np.float32)183 184        results.append(y)185    return results186 187 188@torch.inference_mode()189def numpy2pytorch(imgs):190    h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.0 - 1.0  # so that 127 must be strictly 0.0191    h = h.movedim(-1, 1)192    return h193 194 195def resize_and_center_crop(image, target_width, target_height):196    pil_image = Image.fromarray(image)197    original_width, original_height = pil_image.size198    scale_factor = max(target_width / original_width, target_height / original_height)199    resized_width = int(round(original_width * scale_factor))200    resized_height = int(round(original_height * scale_factor))201    resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)202    left = (resized_width - target_width) / 2203    top = (resized_height - target_height) / 2204    right = (resized_width + target_width) / 2205    bottom = (resized_height + target_height) / 2206    cropped_image = resized_image.crop((left, top, right, bottom))207    return np.array(cropped_image)208 209 210def resize_without_crop(image, target_width, target_height):211    pil_image = Image.fromarray(image)212    resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)213    return np.array(resized_image)214 215 216@torch.inference_mode()217def run_rmbg(img, sigma=0.0):218    H, W, C = img.shape219    assert C == 3220    k = (256.0 / float(H * W)) ** 0.5221    feed = resize_without_crop(img, int(64 * round(W * k)), int(64 * round(H * k)))222    feed = numpy2pytorch([feed]).to(device=device, dtype=torch.float32)223    alpha = rmbg(feed)[0][0]224    alpha = torch.nn.functional.interpolate(alpha, size=(H, W), mode="bilinear")225    alpha = alpha.movedim(1, -1)[0]226    alpha = alpha.detach().float().cpu().numpy().clip(0, 1)227    result = 127 + (img.astype(np.float32) - 127 + sigma) * alpha228    return result.clip(0, 255).astype(np.uint8), alpha229 230 231@torch.inference_mode()232def process(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source):233    bg_source = BGSource(bg_source)234    input_bg = None235 236    if bg_source == BGSource.NONE:237        pass238    elif bg_source == BGSource.LEFT:239        gradient = np.linspace(255, 0, image_width)240        image = np.tile(gradient, (image_height, 1))241        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)242    elif bg_source == BGSource.RIGHT:243        gradient = np.linspace(0, 255, image_width)244        image = np.tile(gradient, (image_height, 1))245        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)246    elif bg_source == BGSource.TOP:247        gradient = np.linspace(255, 0, image_height)[:, None]248        image = np.tile(gradient, (1, image_width))249        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)250    elif bg_source == BGSource.BOTTOM:251        gradient = np.linspace(0, 255, image_height)[:, None]252        image = np.tile(gradient, (1, image_width))253        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)254    else:255        raise 'Wrong initial latent!'256 257    rng = torch.Generator(device=device).manual_seed(int(seed))258 259    fg = resize_and_center_crop(input_fg, image_width, image_height)260 261    concat_conds = numpy2pytorch([fg]).to(device=vae.device, dtype=vae.dtype)262    concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor263 264    conds, unconds = encode_prompt_pair(positive_prompt=prompt + ', ' + a_prompt, negative_prompt=n_prompt)265 266    if input_bg is None:267        latents = t2i_pipe(268            prompt_embeds=conds,269            negative_prompt_embeds=unconds,270            width=image_width,271            height=image_height,272            num_inference_steps=steps,273            num_images_per_prompt=num_samples,274            generator=rng,275            output_type='latent',276            guidance_scale=cfg,277            cross_attention_kwargs={'concat_conds': concat_conds},278        ).images.to(vae.dtype) / vae.config.scaling_factor279    else:280        bg = resize_and_center_crop(input_bg, image_width, image_height)281        bg_latent = numpy2pytorch([bg]).to(device=vae.device, dtype=vae.dtype)282        bg_latent = vae.encode(bg_latent).latent_dist.mode() * vae.config.scaling_factor283        latents = i2i_pipe(284            image=bg_latent,285            strength=lowres_denoise,286            prompt_embeds=conds,287            negative_prompt_embeds=unconds,288            width=image_width,289            height=image_height,290            num_inference_steps=int(round(steps / lowres_denoise)),291            num_images_per_prompt=num_samples,292            generator=rng,293            output_type='latent',294            guidance_scale=cfg,295            cross_attention_kwargs={'concat_conds': concat_conds},296        ).images.to(vae.dtype) / vae.config.scaling_factor297 298    pixels = vae.decode(latents).sample299    pixels = pytorch2numpy(pixels)300    pixels = [resize_without_crop(301        image=p,302        target_width=int(round(image_width * highres_scale / 64.0) * 64),303        target_height=int(round(image_height * highres_scale / 64.0) * 64))304    for p in pixels]305 306    pixels = numpy2pytorch(pixels).to(device=vae.device, dtype=vae.dtype)307    latents = vae.encode(pixels).latent_dist.mode() * vae.config.scaling_factor308    latents = latents.to(device=unet.device, dtype=unet.dtype)309 310    image_height, image_width = latents.shape[2] * 8, latents.shape[3] * 8311 312    fg = resize_and_center_crop(input_fg, image_width, image_height)313    concat_conds = numpy2pytorch([fg]).to(device=vae.device, dtype=vae.dtype)314    concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor315 316    latents = i2i_pipe(317        image=latents,318        strength=highres_denoise,319        prompt_embeds=conds,320        negative_prompt_embeds=unconds,321        width=image_width,322        height=image_height,323        num_inference_steps=int(round(steps / highres_denoise)),324        num_images_per_prompt=num_samples,325        generator=rng,326        output_type='latent',327        guidance_scale=cfg,328        cross_attention_kwargs={'concat_conds': concat_conds},329    ).images.to(vae.dtype) / vae.config.scaling_factor330 331    pixels = vae.decode(latents).sample332 333    return pytorch2numpy(pixels)334 335 336@spaces.GPU337@torch.inference_mode()338def process_relight(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source):339    input_fg, matting = run_rmbg(input_fg)340    results = process(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source)341    return input_fg, results342 343 344quick_prompts = [345    'sunshine from window',346    'neon light, city',347    'sunset over sea',348    'golden time',349    'sci-fi RGB glowing, cyberpunk',350    'natural lighting',351    'warm atmosphere, at home, bedroom',352    'magic lit',353    'evil, gothic, Yharnam',354    'light and shadow',355    'shadow from window',356    'soft studio lighting',357    'home atmosphere, cozy bedroom illumination',358    'neon, Wong Kar-wai, warm'359]360quick_prompts = [[x] for x in quick_prompts]361 362 363quick_subjects = [364    'beautiful woman, detailed face',365    'handsome man, detailed face',366]367quick_subjects = [[x] for x in quick_subjects]368 369 370class BGSource(Enum):371    NONE = "None"372    LEFT = "Left Light"373    RIGHT = "Right Light"374    TOP = "Top Light"375    BOTTOM = "Bottom Light"376 377 378block = gr.Blocks().queue()379with block:380    with gr.Row():381        gr.Markdown("## IC-Light (Relighting with Foreground Condition)")382    with gr.Row():383        gr.Markdown("See also https://github.com/lllyasviel/IC-Light for background-conditioned model and normal estimation")384    with gr.Row():385        with gr.Column():386            with gr.Row():387                input_fg = gr.Image(sources='upload', type="numpy", label="Image", height=480)388                output_bg = gr.Image(type="numpy", label="Preprocessed Foreground", height=480)389            prompt = gr.Textbox(label="Prompt")390            bg_source = gr.Radio(choices=[e.value for e in BGSource],391                                 value=BGSource.NONE.value,392                                 label="Lighting Preference (Initial Latent)", type='value')393            example_quick_subjects = gr.Dataset(samples=quick_subjects, label='Subject Quick List', samples_per_page=1000, components=[prompt])394            example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Lighting Quick List', samples_per_page=1000, components=[prompt])395            relight_button = gr.Button(value="Relight")396 397            with gr.Group():398                with gr.Row():399                    num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)400                    seed = gr.Number(label="Seed", value=12345, precision=0)401 402                with gr.Row():403                    image_width = gr.Slider(label="Image Width", minimum=256, maximum=1024, value=512, step=64)404                    image_height = gr.Slider(label="Image Height", minimum=256, maximum=1024, value=640, step=64)405 406            with gr.Accordion("Advanced options", open=False):407                steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1)408                cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=32.0, value=2, step=0.01)409                lowres_denoise = gr.Slider(label="Lowres Denoise (for initial latent)", minimum=0.1, maximum=1.0, value=0.9, step=0.01)410                highres_scale = gr.Slider(label="Highres Scale", minimum=1.0, maximum=3.0, value=1.5, step=0.01)411                highres_denoise = gr.Slider(label="Highres Denoise", minimum=0.1, maximum=1.0, value=0.5, step=0.01)412                a_prompt = gr.Textbox(label="Added Prompt", value='best quality')413                n_prompt = gr.Textbox(label="Negative Prompt", value='lowres, bad anatomy, bad hands, cropped, worst quality')414        with gr.Column():415            result_gallery = gr.Gallery(height=832, object_fit='contain', label='Outputs')416    with gr.Row():417        dummy_image_for_outputs = gr.Image(visible=False, label='Result')418        gr.Examples(419            fn=lambda *args: [[args[-1]], "imgs/dummy.png"],420            examples=db_examples.foreground_conditioned_examples,421            inputs=[422                input_fg, prompt, bg_source, image_width, image_height, seed, dummy_image_for_outputs423            ],424            outputs=[result_gallery, output_bg],425            run_on_click=True, examples_per_page=1024426        )427    ips = [input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source]428    relight_button.click(fn=process_relight, inputs=ips, outputs=[output_bg, result_gallery])429    example_quick_prompts.click(lambda x, y: ', '.join(y.split(', ')[:2] + [x[0]]), inputs=[example_quick_prompts, prompt], outputs=prompt, show_progress=False, queue=False)430    example_quick_subjects.click(lambda x: x[0], inputs=example_quick_subjects, outputs=prompt, show_progress=False, queue=False)431 432 433block.launch(server_name='0.0.0.0')434