SubstanceSHIFT/SeedVR2-3B

1# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates2# //3# // Licensed under the Apache License, Version 2.0 (the "License");4# // you may not use this file except in compliance with the License.5# // You may obtain a copy of the License at6# //7# //     http://www.apache.org/licenses/LICENSE-2.08# //9# // Unless required by applicable law or agreed to in writing, software10# // distributed under the License is distributed on an "AS IS" BASIS,11# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# // See the License for the specific language governing permissions and13# // limitations under the License.14 15import os16import torch17import mediapy18from einops import rearrange19from omegaconf import OmegaConf20print(os.getcwd())21import datetime22from tqdm import tqdm23from models.dit import na24import gc25 26from data.image.transforms.divisible_crop import DivisibleCrop27from data.image.transforms.na_resize import NaResize28from data.video.transforms.rearrange import Rearrange29if os.path.exists("./projects/video_diffusion_sr/color_fix.py"):30    from projects.video_diffusion_sr.color_fix import wavelet_reconstruction31    use_colorfix=True32else:33    use_colorfix = False34    print('Note!!!!!! Color fix is not avaliable!')35from torchvision.transforms import Compose, Lambda, Normalize36from torchvision.io.video import read_video37import argparse38 39 40from common.distributed import (41    get_device,42    init_torch,43)44 45from common.distributed.advanced import (46    get_data_parallel_rank,47    get_data_parallel_world_size,48    get_sequence_parallel_rank,49    get_sequence_parallel_world_size,50    init_sequence_parallel,51)52 53from projects.video_diffusion_sr.infer import VideoDiffusionInfer54from common.config import load_config55from common.distributed.ops import sync_data56from common.seed import set_seed57from common.partition import partition_by_groups, partition_by_size58 59 60def configure_sequence_parallel(sp_size):61    if sp_size > 1:62        init_sequence_parallel(sp_size)63 64def configure_runner(sp_size):65    config_path = os.path.join('./configs_7b', 'main.yaml')66    config = load_config(config_path)67    runner = VideoDiffusionInfer(config)68    OmegaConf.set_readonly(runner.config, False)69    70    init_torch(cudnn_benchmark=False, timeout=datetime.timedelta(seconds=3600))71    configure_sequence_parallel(sp_size)72    runner.configure_dit_model(device="cuda", checkpoint='./ckpts/seedvr_ema_7b.pth')73    runner.configure_vae_model()74    # Set memory limit.75    if hasattr(runner.vae, "set_memory_limit"):76        runner.vae.set_memory_limit(**runner.config.vae.memory_limit)77    return runner78 79def generation_step(runner, text_embeds_dict, cond_latents):80    def _move_to_cuda(x):81        return [i.to(get_device()) for i in x]82 83    noises = [torch.randn_like(latent) for latent in cond_latents]84    aug_noises = [torch.randn_like(latent) for latent in cond_latents]85    print(f"Generating with noise shape: {noises[0].size()}.")86    noises, aug_noises, cond_latents = sync_data((noises, aug_noises, cond_latents), 0)87    noises, aug_noises, cond_latents = list(88        map(lambda x: _move_to_cuda(x), (noises, aug_noises, cond_latents))89    )90    cond_noise_scale = 0.191 92    def _add_noise(x, aug_noise):93        t = (94            torch.tensor([1000.0], device=get_device())95            * cond_noise_scale96        )97        shape = torch.tensor(x.shape[1:], device=get_device())[None]98        t = runner.timestep_transform(t, shape)99        print(100            f"Timestep shifting from"101            f" {1000.0 * cond_noise_scale} to {t}."102        )103        x = runner.schedule.forward(x, aug_noise, t)104        return x105 106    conditions = [107        runner.get_condition(108            noise,109            task="sr",110            latent_blur=_add_noise(latent_blur, aug_noise),111        )112        for noise, aug_noise, latent_blur in zip(noises, aug_noises, cond_latents)113    ]114 115    with torch.no_grad(), torch.autocast("cuda", torch.bfloat16, enabled=True):116        video_tensors = runner.inference(117            noises=noises,118            conditions=conditions,119            dit_offload=True,120            **text_embeds_dict,121        )122 123    samples = [124        (125            rearrange(video[:, None], "c t h w -> t c h w")126            if video.ndim == 3127            else rearrange(video, "c t h w -> t c h w")128        )129        for video in video_tensors130    ]131    del video_tensors132 133    return samples134 135def generation_loop(runner, video_path='./test_videos', output_dir='./results', batch_size=1, cfg_scale=6.5, cfg_rescale=0.0, sample_steps=50, seed=666, res_h=1280, res_w=720, sp_size=1):136 137    def _build_pos_and_neg_prompt():138        # read positive prompt139        positive_text = "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, \140        hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, \141        skin pore detailing, hyper sharpness, perfect without deformations."142        # read negative prompt143        negative_text = "painting, oil painting, illustration, drawing, art, sketch, oil painting, cartoon, \144        CG Style, 3D render, unreal engine, blurring, dirty, messy, worst quality, low quality, frames, watermark, \145        signature, jpeg artifacts, deformed, lowres, over-smooth"146        return positive_text, negative_text147 148    def _build_test_prompts(video_path):149        positive_text, negative_text = _build_pos_and_neg_prompt()150        original_videos = []151        prompts = {}152        video_list = os.listdir(video_path)153        for f in video_list:154            if f.endswith(".mp4"):155                original_videos.append(f)156                prompts[f] = positive_text157        print(f"Total prompts to be generated: {len(original_videos)}")158        return original_videos, prompts, negative_text159 160    def _extract_text_embeds():161        # Text encoder forward.162        positive_prompts_embeds = []163        for texts_pos in tqdm(original_videos_local):164            text_pos_embeds = torch.load('pos_emb.pt')165            text_neg_embeds = torch.load('neg_emb.pt')166 167            positive_prompts_embeds.append(168                {"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]}169            )170        gc.collect()171        torch.cuda.empty_cache()172        return positive_prompts_embeds173 174    def cut_videos(videos, sp_size):175        t = videos.size(1)176        if t <= 4 * sp_size:177            print(f"Cut input video size: {videos.size()}")178            padding = [videos[:, -1].unsqueeze(1)] * (4 * sp_size - t + 1)179            padding = torch.cat(padding, dim=1)180            videos = torch.cat([videos, padding], dim=1)181            return videos182        if (t - 1) % (4 * sp_size) == 0:183            return videos184        else:185            padding = [videos[:, -1].unsqueeze(1)] * (186                4 * sp_size - ((t - 1) % (4 * sp_size))187            )188            padding = torch.cat(padding, dim=1)189            videos = torch.cat([videos, padding], dim=1)190            assert (videos.size(1) - 1) % (4 * sp_size) == 0191            return videos192 193    # classifier-free guidance194    runner.config.diffusion.cfg.scale = cfg_scale195    runner.config.diffusion.cfg.rescale = cfg_rescale196    # sampling steps197    runner.config.diffusion.timesteps.sampling.steps = sample_steps198    runner.configure_diffusion()199 200    # set random seed201    set_seed(seed, same_across_ranks=True)202    os.makedirs(output_dir, exist_ok=True)203    tgt_path = output_dir204 205    # get test prompts206    original_videos, _, _ = _build_test_prompts(video_path)207 208    # divide the prompts into different groups209    original_videos_group = partition_by_groups(210        original_videos,211        get_data_parallel_world_size() // get_sequence_parallel_world_size(),212    )213    # store prompt mapping214    original_videos_local = original_videos_group[215        get_data_parallel_rank() // get_sequence_parallel_world_size()216    ]217    original_videos_local = partition_by_size(original_videos_local, batch_size)218 219    # pre-extract the text embeddings220    positive_prompts_embeds = _extract_text_embeds()221 222    video_transform = Compose(223        [224            NaResize(225                resolution=(226                    res_h * res_w227                )228                ** 0.5,229                mode="area",230                # Upsample image, model only trained for high res.231                downsample_only=False,232            ),233            Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),234            DivisibleCrop((16, 16)),235            Normalize(0.5, 0.5),236            Rearrange("t c h w -> c t h w"),237        ]238    )239 240    # generation loop241    for videos, text_embeds in tqdm(zip(original_videos_local, positive_prompts_embeds)):242        # read condition latents243        cond_latents = []244        for video in videos:245            video = (246                read_video(247                   os.path.join(video_path, video), output_format="TCHW"248                )[0]249                / 255.0250            )251            print(f"Read video size: {video.size()}")252            cond_latents.append(video_transform(video.to(get_device())))253 254        ori_lengths = [video.size(1) for video in cond_latents]255        input_videos = cond_latents256        cond_latents = [cut_videos(video, sp_size) for video in cond_latents]257 258        runner.dit.to("cpu")259        print(f"Encoding videos: {list(map(lambda x: x.size(), cond_latents))}")260        runner.vae.to(get_device())261        cond_latents = runner.vae_encode(cond_latents)262        runner.vae.to("cpu")263        runner.dit.to(get_device())264 265        for i, emb in enumerate(text_embeds["texts_pos"]):266            text_embeds["texts_pos"][i] = emb.to(get_device())267        for i, emb in enumerate(text_embeds["texts_neg"]):268            text_embeds["texts_neg"][i] = emb.to(get_device())269 270        samples = generation_step(runner, text_embeds, cond_latents=cond_latents)271        runner.dit.to("cpu")272        del cond_latents273 274        # dump samples to the output directory275        if get_sequence_parallel_rank() == 0:276            for path, input, sample, ori_length in zip(277                videos, input_videos, samples, ori_lengths278            ):279                if ori_length < sample.shape[0]:280                    sample = sample[:ori_length]281                filename = os.path.join(tgt_path, os.path.basename(path))282                # color fix283                input = (284                    rearrange(input[:, None], "c t h w -> t c h w")285                    if input.ndim == 3286                    else rearrange(input, "c t h w -> t c h w")287                )288                if use_colorfix:289                    sample = wavelet_reconstruction(290                        sample.to("cpu"), input[: sample.size(0)].to("cpu")291                    )292                else:293                    sample = sample.to("cpu")294                sample = (295                    rearrange(sample[:, None], "t c h w -> t h w c")296                    if sample.ndim == 3297                    else rearrange(sample, "t c h w -> t h w c")298                )299                sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round()300                sample = sample.to(torch.uint8).numpy()301 302                if sample.shape[0] == 1:303                    mediapy.write_image(filename, sample.squeeze(0))304                else:305                    mediapy.write_video(306                        filename, sample, fps=24307                    )308        gc.collect()309        torch.cuda.empty_cache()310 311if __name__ == "__main__":312    parser = argparse.ArgumentParser() 313    parser.add_argument("--video_path", type=str, default="./test_videos")314    parser.add_argument("--output_dir", type=str, default="./results")315    parser.add_argument("--cfg_scale", type=float, default=6.5)316    parser.add_argument("--sample_steps", type=int, default=50)317    parser.add_argument("--seed", type=int, default=666)318    parser.add_argument("--res_h", type=int, default=720)319    parser.add_argument("--res_w", type=int, default=1280)320    parser.add_argument("--sp_size", type=int, default=1)321    args = parser.parse_args()322 323    runner = configure_runner(args.sp_size)324    generation_loop(runner, **vars(args))325