Update app.py

app.py

@@ -14,73 +14,65 @@
 import spaces
 import subprocess
 import os
-import sys # <-- ADICIONADO PARA MANIPULAR O CAMINHO DO PYTHON
+import sys
 
-# Clone the repository to ensure all files are available
-# Make sure git-lfs is installed
+# --- Setup: Clone repository and add it to Python Path ---
+# This section ensures all necessary code and model files are available.
+
+# 1. Clone the repository with all its files
 subprocess.run("git lfs install", shell=True, check=True)
-# Clone the repository only if it doesn't exist
 if not os.path.exists("SeedVR2-3B"):
+    print("Cloning SeedVR2-3B repository...")
     subprocess.run("git clone https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B", shell=True, check=True)
 
-# Define the repository directory
-repo_dir = 'SeedVR2-3B'
-# Change the current working directory to the cloned repository
-os.chdir(repo_dir)
-# Add the repository directory to the Python path to allow imports
-sys.path.insert(0, os.path.abspath('.')) # <-- CORREÇÃO PRINCIPAL AQUI
+# 2. Add the cloned repository's directory to Python's module search path
+repo_dir = "SeedVR2-3B"
+# This allows us to import modules like 'data', 'common', etc., from the cloned repo.
+sys.path.insert(0, os.path.abspath(repo_dir))
+print(f"Repository directory '{os.path.abspath(repo_dir)}' added to Python path.")
+
+# --- Main Application Code ---
+# All file paths will now be relative to the cloned repository directory.
 
 import torch
 import mediapy
 from einops import rearrange
 from omegaconf import OmegaConf
-print(os.getcwd())
 import datetime
 from tqdm import tqdm
 import gc
-
-from data.image.transforms.divisible_crop import DivisibleCrop
-from data.image.transforms.na_resize import NaResize
-from data.video.transforms.rearrange import Rearrange
-if os.path.exists("./projects/video_diffusion_sr/color_fix.py"):
-    from projects.video_diffusion_sr.color_fix import wavelet_reconstruction
-    use_colorfix=True
-else:
-    use_colorfix = False
-    print('Note!!!!!! Color fix is not avaliable!')
-from torchvision.transforms import Compose, Lambda, Normalize
-from torchvision.io.video import read_video
-import argparse
 from PIL import Image
-
-from common.distributed import (
-    get_device,
-    init_torch,
-)
-
-from common.distributed.advanced import (
-    get_data_parallel_rank,
-    get_data_parallel_world_size,
-    get_sequence_parallel_rank,
-    get_sequence_parallel_world_size,
-    init_sequence_parallel,
-)
-
-from projects.video_diffusion_sr.infer import VideoDiffusionInfer
-from common.config import load_config
-from common.distributed.ops import sync_data
-from common.seed import set_seed
-from common.partition import partition_by_groups, partition_by_size
-
 import gradio as gr
 from pathlib import Path
-from urllib.parse import urlparse
-from torch.hub import download_url_to_file, get_dir
 import shlex
 import uuid
 import mimetypes
 import torchvision.transforms as T
+from torchvision.transforms import Compose, Lambda, Normalize
+from torchvision.io.video import read_video
 
+# Imports from the cloned repository
+from data.image.transforms.divisible_crop import DivisibleCrop
+from data.image.transforms.na_resize import NaResize
+from data.video.transforms.rearrange import Rearrange
+from common.config import load_config
+from common.distributed import init_torch
+from common.distributed.advanced import init_sequence_parallel
+from common.seed import set_seed
+from common.partition import partition_by_size
+from projects.video_diffusion_sr.infer import VideoDiffusionInfer
+from common.distributed.ops import sync_data
+
+# Check for color_fix utility
+color_fix_path = os.path.join(repo_dir, "projects/video_diffusion_sr/color_fix.py")
+if os.path.exists(color_fix_path):
+    from projects.video_diffusion_sr.color_fix import wavelet_reconstruction
+    use_colorfix = True
+else:
+    use_colorfix = False
+    print('Note!!!!!! Color fix is not available!')
+
+# --- Environment and Dependencies Setup ---
 os.environ["MASTER_ADDR"] = "127.0.0.1"
 os.environ["MASTER_PORT"] = "12355"
 os.environ["RANK"] = str(0)
@@ -92,33 +84,34 @@ subprocess.run(
     shell=True,
 )
 
-# Install apex from the local wheel file
-if os.path.exists("apex-0.1-cp310-cp310-linux_x86_64.whl"):
-    subprocess.run(shlex.split("pip install apex-0.1-cp310-cp310-linux_x86_64.whl"))
-    print(f"✅ setup completed Apex")
+apex_wheel_path = os.path.join(repo_dir, "apex-0.1-cp310-cp310-linux_x86_64.whl")
+if os.path.exists(apex_wheel_path):
+    subprocess.run(shlex.split(f"pip install {apex_wheel_path}"))
+    print("✅ Apex setup completed.")
 
+# --- Core Functions ---
 
 def configure_sequence_parallel(sp_size):
     if sp_size > 1:
         init_sequence_parallel(sp_size)
 
-
 def configure_runner(sp_size):
-    config_path = os.path.join('./configs_3b', 'main.yaml')
+    config_path = os.path.join(repo_dir, 'configs_3b', 'main.yaml')
+    checkpoint_path = os.path.join(repo_dir, 'ckpts', 'seedvr2_ema_3b.pth')
+    
     config = load_config(config_path)
     runner = VideoDiffusionInfer(config)
     OmegaConf.set_readonly(runner.config, False)
     
     init_torch(cudnn_benchmark=False, timeout=datetime.timedelta(seconds=3600))
     configure_sequence_parallel(sp_size)
-    runner.configure_dit_model(device="cuda", checkpoint='./ckpts/seedvr2_ema_3b.pth')
+    runner.configure_dit_model(device="cuda", checkpoint=checkpoint_path)
     runner.configure_vae_model()
-    # Set memory limit.
+    
     if hasattr(runner.vae, "set_memory_limit"):
         runner.vae.set_memory_limit(**runner.config.vae.memory_limit)
     return runner
 
-
 def generation_step(runner, text_embeds_dict, cond_latents):
     def _move_to_cuda(x):
         return [i.to(torch.device("cuda")) for i in x]
@@ -127,68 +120,44 @@ def generation_step(runner, text_embeds_dict, cond_latents):
     aug_noises = [torch.randn_like(latent) for latent in cond_latents]
     print(f"Generating with noise shape: {noises[0].size()}.")
     noises, aug_noises, cond_latents = sync_data((noises, aug_noises, cond_latents), 0)
-    noises, aug_noises, cond_latents = list(
-        map(lambda x: _move_to_cuda(x), (noises, aug_noises, cond_latents))
-    )
+    noises, aug_noises, cond_latents = list(map(_move_to_cuda, (noises, aug_noises, cond_latents)))
     cond_noise_scale = 0.1
 
     def _add_noise(x, aug_noise):
-        t = (
-            torch.tensor([1000.0], device=torch.device("cuda"))
-            * cond_noise_scale
-        )
+        t = torch.tensor([1000.0], device=torch.device("cuda")) * cond_noise_scale
         shape = torch.tensor(x.shape[1:], device=torch.device("cuda"))[None]
         t = runner.timestep_transform(t, shape)
-        print(
-            f"Timestep shifting from"
-            f" {1000.0 * cond_noise_scale} to {t}."
-        )
+        print(f"Timestep shifting from {1000.0 * cond_noise_scale} to {t}.")
         x = runner.schedule.forward(x, aug_noise, t)
         return x
 
     conditions = [
-        runner.get_condition(
-            noise,
-            task="sr",
-            latent_blur=_add_noise(latent_blur, aug_noise),
-        )
+        runner.get_condition(noise, task="sr", latent_blur=_add_noise(latent_blur, aug_noise))
         for noise, aug_noise, latent_blur in zip(noises, aug_noises, cond_latents)
     ]
 
     with torch.no_grad(), torch.autocast("cuda", torch.bfloat16, enabled=True):
         video_tensors = runner.inference(
-            noises=noises,
-            conditions=conditions,
-            dit_offload=False,
-            **text_embeds_dict,
+            noises=noises, conditions=conditions, dit_offload=False, **text_embeds_dict
         )
 
-    samples = [
-        (
-            rearrange(video[:, None], "c t h w -> t c h w")
-            if video.ndim == 3
-            else rearrange(video, "c t h w -> t c h w")
-        )
-        for video in video_tensors
-    ]
+    samples = [rearrange(video, "c t h w -> t c h w") for video in video_tensors]
     del video_tensors
-
     return samples
 
-
-def generation_loop(video_path='./test_videos', seed=666, fps_out=12, batch_size=1, cfg_scale=1.0, cfg_rescale=0.0, sample_steps=1, res_h=1280, res_w=720, sp_size=1):
+@spaces.GPU
+def generation_loop(video_path, seed=666, fps_out=24, batch_size=1, cfg_scale=1.0, cfg_rescale=0.0, sample_steps=1, res_h=1280, res_w=720, sp_size=1):
+    if video_path is None:
+        return None, None, None
+        
     runner = configure_runner(1)
 
     def _extract_text_embeds():
-        # Text encoder forward.
         positive_prompts_embeds = []
-        for texts_pos in tqdm(original_videos_local):
-            text_pos_embeds = torch.load('pos_emb.pt')
-            text_neg_embeds = torch.load('neg_emb.pt')
-
-            positive_prompts_embeds.append(
-                {"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]}
-            )
+        for _ in original_videos_local:
+            text_pos_embeds = torch.load(os.path.join(repo_dir, 'pos_emb.pt'))
+            text_neg_embeds = torch.load(os.path.join(repo_dir, 'neg_emb.pt'))
+            positive_prompts_embeds.append({"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]})
         gc.collect()
         torch.cuda.empty_cache()
         return positive_prompts_embeds
@@ -198,239 +167,56 @@ def generation_loop(video_path='./test_videos', seed=666, fps_out=12, batch_size
             videos = videos[:, :121]
         t = videos.size(1)
         if t <= 4 * sp_size:
-            print(f"Cut input video size: {videos.size()}")
-            padding = [videos[:, -1].unsqueeze(1)] * (4 * sp_size - t + 1)
-            padding = torch.cat(padding, dim=1)
-            videos = torch.cat([videos, padding], dim=1)
+            padding_needed = 4 * sp_size - t + 1
+            if padding_needed > 0:
+                padding = torch.cat([videos[:, -1].unsqueeze(1)] * padding_needed, dim=1)
+                videos = torch.cat([videos, padding], dim=1)
             return videos
         if (t - 1) % (4 * sp_size) == 0:
             return videos
         else:
-            padding = [videos[:, -1].unsqueeze(1)] * (
-                4 * sp_size - ((t - 1) % (4 * sp_size))
-            )
-            padding = torch.cat(padding, dim=1)
+            padding_needed = 4 * sp_size - ((t - 1) % (4 * sp_size))
+            padding = torch.cat([videos[:, -1].unsqueeze(1)] * padding_needed, dim=1)
             videos = torch.cat([videos, padding], dim=1)
             assert (videos.size(1) - 1) % (4 * sp_size) == 0
             return videos
 
-    # classifier-free guidance
     runner.config.diffusion.cfg.scale = cfg_scale
     runner.config.diffusion.cfg.rescale = cfg_rescale
-    # sampling steps
     runner.config.diffusion.timesteps.sampling.steps = sample_steps
     runner.configure_diffusion()
 
-    # set random seed
-    seed = seed % (2**32) # avoid over range
+    seed = int(seed) % (2**32)
     set_seed(seed, same_across_ranks=True)
-    os.makedirs('output/', exist_ok=True)
+    output_base_dir = "output"
+    os.makedirs(output_base_dir, exist_ok=True)
 
-    # get test prompts
     original_videos = [os.path.basename(video_path)]
-
-    # divide the prompts into different groups
-    original_videos_group = original_videos
-    # store prompt mapping
-    original_videos_local = original_videos_group
-    original_videos_local = partition_by_size(original_videos_local, batch_size)
-
-    # pre-extract the text embeddings
+    original_videos_local = partition_by_size(original_videos, batch_size)
     positive_prompts_embeds = _extract_text_embeds()
 
-    video_transform = Compose(
-        [
-            NaResize(
-                resolution=(
-                    res_h * res_w
-                )
-                ** 0.5,
-                mode="area",
-                # Upsample image, model only trained for high res.
-                downsample_only=False,
-            ),
-            Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
-            DivisibleCrop((16, 16)),
-            Normalize(0.5, 0.5),
-            Rearrange("t c h w -> c t h w"),
-        ]
-    )
+    video_transform = Compose([
+        NaResize(resolution=(res_h * res_w) ** 0.5, mode="area", downsample_only=False),
+        Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
+        DivisibleCrop((16, 16)),
+        Normalize(0.5, 0.5),
+        Rearrange("t c h w -> c t h w"),
+    ])
 
-    # generation loop
     for videos, text_embeds in tqdm(zip(original_videos_local, positive_prompts_embeds)):
-        # read condition latents
         cond_latents = []
-        for video in videos:
+        for _ in videos:
             media_type, _ = mimetypes.guess_type(video_path)
             is_image = media_type and media_type.startswith("image")
             is_video = media_type and media_type.startswith("video")
+            
             if is_video:
-                video = (
-                    read_video(
-                       video_path, output_format="TCHW"
-                    )[0]
-                    / 255.0
-                )
+                video, _, _ = read_video(video_path, output_format="TCHW")
+                video = video / 255.0
                 if video.size(0) > 121:
                     video = video[:121]
                 print(f"Read video size: {video.size()}")
-                output_dir = 'output/' + str(uuid.uuid4()) + '.mp4'
-            else:
+                output_dir = os.path.join(output_base_dir, f"{uuid.uuid4()}.mp4")
+            elif is_image:
                 img = Image.open(video_path).convert("RGB")
-                img_tensor = T.ToTensor()(img).unsqueeze(0)  # (1, C, H, W)
-                video = img_tensor.permute(0, 1, 2, 3)  # (T=1, C, H, W)
-                print(f"Read Image size: {video.size()}")
-                output_dir = 'output/' + str(uuid.uuid4()) + '.png'
-            cond_latents.append(video_transform(video.to(torch.device("cuda"))))
-
-        ori_lengths = [video.size(1) for video in cond_latents]
-        input_videos = cond_latents
-        if is_video:
-            cond_latents = [cut_videos(video, sp_size) for video in cond_latents]
-
-        print(f"Encoding videos: {list(map(lambda x: x.size(), cond_latents))}")
-        cond_latents = runner.vae_encode(cond_latents)
-
-        for i, emb in enumerate(text_embeds["texts_pos"]):
-            text_embeds["texts_pos"][i] = emb.to(torch.device("cuda"))
-        for i, emb in enumerate(text_embeds["texts_neg"]):
-            text_embeds["texts_neg"][i] = emb.to(torch.device("cuda"))
-
-        samples = generation_step(runner, text_embeds, cond_latents=cond_latents)
-        del cond_latents
-
-        # dump samples to the output directory
-        for path, input, sample, ori_length in zip(
-            videos, input_videos, samples, ori_lengths
-        ):
-            if ori_length < sample.shape[0]:
-                sample = sample[:ori_length]
-            # color fix
-            input = (
-                rearrange(input[:, None], "c t h w -> t c h w")
-                if input.ndim == 3
-                else rearrange(input, "c t h w -> t c h w")
-            )
-            if use_colorfix:
-                sample = wavelet_reconstruction(
-                    sample.to("cpu"), input[: sample.size(0)].to("cpu")
-                )
-            else:
-                sample = sample.to("cpu")
-            sample = (
-                rearrange(sample[:, None], "t c h w -> t h w c")
-                if sample.ndim == 3
-                else rearrange(sample, "t c h w -> t h w c")
-            )
-            sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round()
-            sample = sample.to(torch.uint8).numpy()
-
-            if is_image:
-                mediapy.write_image(output_dir, sample[0])
-            else:
-                mediapy.write_video(
-                    output_dir, sample, fps=fps_out
-                )
-
-        gc.collect()
-        torch.cuda.empty_cache()
-        if is_image:
-            return output_dir, None, output_dir
-        else:
-            return None, output_dir, output_dir
-
-
-with gr.Blocks(title="SeedVR2: One-Step Video Restoration via Diffusion Adversarial Post-Training") as demo:
-    # Top logo and title
-    gr.HTML("""
-        <div style='text-align:center; margin-bottom: 10px;'>
-            <img src='assets/seedvr_logo.png' style='height:40px;' alt='SeedVR logo'/>
-        </div>
-        <p><b>Official Gradio demo</b> for 
-        <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>
-        <b>SeedVR2: One-Step Video Restoration via Diffusion Adversarial Post-Training</b></a>.<br>
-        🔥 <b>SeedVR2</b> is a one-step image and video restoration algorithm for real-world and AIGC content.
-        </p>
-    """)
-
-    # Interface
-    with gr.Row():
-        input_video = gr.File(label="Upload image or video", type="filepath")
-        seed = gr.Number(label="Seeds", value=666)
-        fps = gr.Number(label="fps", value=24)
-
-    with gr.Row():
-        output_video = gr.Video(label="Output_Video")
-        output_image = gr.Image(label="Output_Image")
-        download_link = gr.File(label="Download the output")
-
-    run_button = gr.Button("Run")
-    run_button.click(fn=generation_loop, inputs=[input_video, seed, fps], outputs=[output_image, output_video, download_link])
-
-    # Examples
-    gr.Examples(
-        examples=[
-            ["01.mp4", 4, 24],
-            ["02.mp4", 4, 24],
-            ["03.mp4", 4, 24],
-        ],
-        inputs=[input_video, seed, fps]
-    )
-
-    # Article/Footer
-    gr.HTML("""
-        <hr>
-        <p>If you find SeedVR helpful, please ⭐ the 
-        <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>GitHub repository</a>:</p>
-
-        <a href="https://github.com/ByteDance-Seed/SeedVR" target="_blank">
-            <img src="https://img.shields.io/github/stars/ByteDance-Seed/SeedVR?style=social" alt="GitHub Stars">
-        </a>
-
-        <h4>Notice</h4>
-        <p>This demo supports up to <b>720p and 121 frames for videos or 2k images</b>.  
-        For other use cases (image restoration beyond 2K, video resolutions beyond 720p, etc), check the <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>GitHub repo</a>.</p>
-
-        <h4>Limitations</h4>
-        <p>May fail on heavy degradations or small-motion AIGC clips, causing oversharpening or poor restoration.</p>
-
-        <h4>Citation</h4>
-        <pre style="font-size: 12px;">
-        @article{wang2025seedvr2,
-            title={SeedVR2: One-Step Video Restoration via Diffusion Adversarial Post-Training},
-            author={Wang, Jianyi and Lin, Shanchuan and Lin, Zhijie and Ren, Yuxi and Wei, Meng and Yue, Zongsheng and Zhou, Shangchen and Chen, Hao and Zhao, Yang and Yang, Ceyuan and Xiao, Xuefeng and Loy, Chen Change and Jiang, Lu},
-            booktitle={arXiv preprint arXiv:2506.05301},
-            year={2025}
-        }
-
-        @inproceedings{wang2025seedvr,
-            title={SeedVR: Seeding Infinity in Diffusion Transformer Towards Generic Video Restoration},
-            author={Wang, Jianyi and Lin, Zhijie and Wei, Meng and Zhao, Yang and Yang, Ceyuan and Loy, Chen Change and Jiang, Lu},
-            booktitle={CVPR},
-            year={2025}
-        }
-        </pre>
-
-        <h4>License</h4>
-        <p>Licensed under the 
-        <a href="http://www.apache.org/licenses/LICENSE-2.0" target="_blank">Apache 2.0 License</a>.</p>
-
-        <h4>Contact</h4>
-        <p>Email: <b>[email protected]</b></p>
-
-        <p>
-        <a href="https://twitter.com/Iceclearwjy">
-        <img src="https://img.shields.io/twitter/follow/Iceclearwjy?label=%40Iceclearwjy&style=social" alt="Twitter Follow">
-        </a>
-        <a href="https://github.com/IceClear">
-        <img src="https://img.shields.io/github/followers/IceClear?style=social" alt="GitHub Follow">
-        </a>
-        </p>
-
-        <p style="text-align:center;">
-        <img src="https://visitor-badge.laobi.icu/badge?page_id=ByteDance-Seed/SeedVR" alt="visitors">
-        </p>
-    """)
-
-demo.queue()
-demo.launch()
+                img_tensor = T.ToTensor()(img).uns