Update app.py

app.py

@@ -78,15 +78,18 @@ os.environ["MASTER_PORT"] = "12355"
 os.environ["RANK"] = str(0)
 os.environ["WORLD_SIZE"] = str(1)
 
+# CORREÇÃO: Usar sys.executable para chamar o pip corretamente
+python_executable = sys.executable
 subprocess.run(
-    "pip install flash-attn --no-build-isolation",
-    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-    shell=True,
+    [python_executable, "-m", "pip", "install", "flash-attn", "--no-build-isolation"],
+    env={**os.environ, "FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    check=True
 )
 
 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}"))
+    # CORREÇÃO: Usar sys.executable aqui também
+    subprocess.run([python_executable, "-m", "pip", "install", apex_wheel_path], check=True)
     print("✅ Apex setup completed.")
 
 # --- Core Functions ---
@@ -219,4 +222,93 @@ def generation_loop(video_path, seed=666, fps_out=24, batch_size=1, cfg_scale=1.
                 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).uns
+                img_tensor = T.ToTensor()(img).unsqueeze(0)
+                video = img_tensor
+                print(f"Read Image size: {video.size()}")
+                output_dir = os.path.join(output_base_dir, f"{uuid.uuid4()}.png")
+            else:
+                raise ValueError("Unsupported file type")
+                
+            cond_latents.append(video_transform(video.to(torch.device("cuda"))))
+
+        ori_lengths = [v.size(1) for v in cond_latents]
+        input_videos = cond_latents
+        if is_video:
+            cond_latents = [cut_videos(v, sp_size) for v in cond_latents]
+
+        print(f"Encoding videos: {[v.size() for v in 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
+
+        for _, input_tensor, sample, ori_length in zip(videos, input_videos, samples, ori_lengths):
+            if ori_length < sample.shape[0]:
+                sample = sample[:ori_length]
+
+            input_tensor = rearrange(input_tensor, "c t h w -> t c h w")
+            if use_colorfix:
+                sample = wavelet_reconstruction(sample.to("cpu"), input_tensor[:sample.size(0)].to("cpu"))
+            else:
+                sample = sample.to("cpu")
+
+            sample = 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
+
+# --- Gradio UI ---
+
+with gr.Blocks(title="SeedVR2: One-Step Video Restoration") as demo:
+    logo_path = os.path.join(repo_dir, "assets/seedvr_logo.png")
+    gr.HTML(f"""
+        <div style='text-align:center; margin-bottom: 10px;'>
+            <img src='file/{logo_path}' 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>
+    """)
+
+    with gr.Row():
+        input_file = gr.File(label="Upload image or video", type="filepath")
+        with gr.Column():
+            seed = gr.Number(label="Seed", value=666)
+            fps = gr.Number(label="Output FPS (for video)", value=24)
+            
+    run_button = gr.Button("Run")
+    
+    with gr.Row():
+        output_image = gr.Image(label="Output Image")
+        output_video = gr.Video(label="Output Video")
+        
+    download_link = gr.File(label="Download the output")
+
+    run_button.click(fn=generation_loop, inputs=[input_file, seed, fps], outputs=[output_image, output_video, download_link])
+
+    gr.HTML("""
+        <hr>
+        <p>If you find SeedVR helpful, please ⭐ the <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>GitHub repository</a>: 
+        <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></p>
+        <h4>Notice</h4>
+        <p>This demo supports up to <b>720p and 121 frames for videos or 2k images</b>. For other use cases, 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>
+    """)
+
+demo.queue().launch(share=True)