Update app.py

app.py

@@ -11,25 +11,25 @@
 # // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # // See the License for the specific language governing permissions and
 # // limitations under the License.
-import os
-import torch
-import mediapy
-import sys
+import spaces
 import subprocess
-from einops import rearrange
-from omegaconf import OmegaConf
-print(os.getcwd())
-import datetime
-from tqdm import tqdm
-import gc
+import os
+import sys # <-- ADICIONADO PARA MANIPULAR O CAMINHO DO PYTHON
 
 # Clone the repository to ensure all files are available
+# Make sure git-lfs is installed
 subprocess.run("git lfs install", shell=True, check=True)
-subprocess.run("git clone https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B", shell=True, check=True)
-os.chdir('SeedVR2-3B')
+# Clone the repository only if it doesn't exist
+if not os.path.exists("SeedVR2-3B"):
+    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
 
-import os
 import torch
 import mediapy
 from einops import rearrange
@@ -92,8 +92,10 @@ subprocess.run(
     shell=True,
 )
 
-subprocess.run(shlex.split("pip install apex-0.1-cp310-cp310-linux_x86_64.whl"))
-print(f"✅ setup completed Apex")
+# 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")
 
 
 def configure_sequence_parallel(sp_size):
@@ -225,7 +227,7 @@ def generation_loop(video_path='./test_videos', seed=666, fps_out=12, batch_size
     os.makedirs('output/', exist_ok=True)
 
     # get test prompts
-    original_videos = [video_path.split('/')[-1]]
+    original_videos = [os.path.basename(video_path)]
 
     # divide the prompts into different groups
     original_videos_group = original_videos
@@ -265,7 +267,7 @@ def generation_loop(video_path='./test_videos', seed=666, fps_out=12, batch_size
             if is_video:
                 video = (
                     read_video(
-                       os.path.join(video_path), output_format="TCHW"
+                       video_path, output_format="TCHW"
                     )[0]
                     / 255.0
                 )
@@ -274,24 +276,6 @@ def generation_loop(video_path='./test_videos', seed=666, fps_out=12, batch_size
                 print(f"Read video size: {video.size()}")
                 output_dir = 'output/' + str(uuid.uuid4()) + '.mp4'
             else:
-                video_transform = Compose(
-                    [
-                        NaResize(
-                            resolution=(
-                                2560 * 1440
-                            )
-                            ** 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"),
-                    ]
-                )
-                
                 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)
@@ -304,12 +288,8 @@ def generation_loop(video_path='./test_videos', seed=666, fps_out=12, batch_size
         if is_video:
             cond_latents = [cut_videos(video, sp_size) for video in cond_latents]
 
-        # runner.dit.to("cpu")
         print(f"Encoding videos: {list(map(lambda x: x.size(), cond_latents))}")
-        # runner.vae.to(torch.device("cuda"))
         cond_latents = runner.vae_encode(cond_latents)
-        # runner.vae.to("cpu")
-        # runner.dit.to(torch.device("cuda"))
 
         for i, emb in enumerate(text_embeds["texts_pos"]):
             text_embeds["texts_pos"][i] = emb.to(torch.device("cuda"))
@@ -317,7 +297,6 @@ def generation_loop(video_path='./test_videos', seed=666, fps_out=12, batch_size
             text_embeds["texts_neg"][i] = emb.to(torch.device("cuda"))
 
         samples = generation_step(runner, text_embeds, cond_latents=cond_latents)
-        # runner.dit.to("cpu")
         del cond_latents
 
         # dump samples to the output directory
@@ -353,7 +332,6 @@ def generation_loop(video_path='./test_videos', seed=666, fps_out=12, batch_size
                     output_dir, sample, fps=fps_out
                 )
 
-        # print(f"Generated video size: {sample.shape}")
         gc.collect()
         torch.cuda.empty_cache()
         if is_image:
@@ -366,7 +344,7 @@ with gr.Blocks(title="SeedVR2: One-Step Video Restoration via Diffusion Adversar
     # Top logo and title
     gr.HTML("""
         <div style='text-align:center; margin-bottom: 10px;'>
-            <img src='https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/assets/seedvr_logo.png' style='height:40px;' alt='SeedVR logo'/>
+            <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'>
@@ -392,9 +370,9 @@ with gr.Blocks(title="SeedVR2: One-Step Video Restoration via Diffusion Adversar
     # Examples
     gr.Examples(
         examples=[
-            ["./01.mp4", 4, 24],
-            ["./02.mp4", 4, 24],
-            ["./03.mp4", 4, 24],
+            ["01.mp4", 4, 24],
+            ["02.mp4", 4, 24],
+            ["03.mp4", 4, 24],
         ],
         inputs=[input_video, seed, fps]
     )