Update app.py

app.py

@@ -322,7 +322,25 @@ 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:
-                img = Image.open(input_file.name).convert("RGB")
+                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)
                 print(f"Read Image size: {video.size()}")
@@ -439,8 +457,8 @@ with gr.Blocks(title="SeedVR2: One-Step Video Restoration via Diffusion Adversar
         </a>
 
         <h4>Notice</h4>
-        <p>This demo supports up to <b>720p</b> and <b>121 frames</b>.  
-        For other use cases (image restoration, video resolutions beyond 720p, etc), check the <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>GitHub repo</a>.</p>
+        <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>