Add parameter tuning

app.py

@@ -11,14 +11,19 @@ examples = [
 
 # Define a GPU-decorated function for enhancement
 @spaces.GPU(duration=120)
-def enhance_with_gpu(input_video, input_text):
+def enhance_with_gpu(input_video, input_text, upscale, max_chunk_len, chunk_size):
+    """在每次调用时创建新的 STAR_sr 实例，确保参数正确传递"""
+    star = STAR_sr(
+        result_dir="./results/",
+        upscale=upscale,
+        max_chunk_len=max_chunk_len,
+        chunk_size=chunk_size
+    )
     return star.enhance_a_video(input_video, input_text)
 
 def star_demo(result_dir="./tmp/"):
     css = """#input_video {max-width: 1024px !important} #output_vid {max-width: 2048px; max-height:1280px}"""
-    global star
-    star = STAR_sr(result_dir)
-    
+
     with gr.Blocks(analytics_enabled=False, css=css) as star_iface:
         gr.Markdown(
             "<div align='center'> <h1> STAR: Spatial-Temporal Augmentation with Text-to-Video Models for Real-World Video Super-Resolution </span> </h1> \
@@ -30,15 +35,16 @@ def star_demo(result_dir="./tmp/"):
             with gr.Column():
                 with gr.Row():
                     with gr.Column():
-                        with gr.Row():
-                            input_video = gr.Video(label="Input Video", elem_id="input_video")
-                        with gr.Row():
-                            input_text = gr.Text(label="Prompts")
+                        input_video = gr.Video(label="Input Video", elem_id="input_video")
+                        input_text = gr.Text(label="Prompts")
+                        
+                        upscale = gr.Slider(1, 4, value=4, step=1, label="Upscale Factor")
+                        max_chunk_len = gr.Slider(1, 32, value=24, step=1, label="Input Chunk Length")
+                        chunk_size = gr.Slider(1, 5, value=3, step=1, label="Decode Chunk Size")
+
                         end_btn = gr.Button("Generate")
-                    with gr.Row():
-                        output_video = gr.Video(
-                            label="Generated Video", elem_id="output_vid", autoplay=True, show_share_button=True
-                        )
+
+                    output_video = gr.Video(label="Generated Video", elem_id="output_vid", autoplay=True, show_share_button=True)
 
                 gr.Examples(
                     examples=examples,
@@ -47,15 +53,15 @@ def star_demo(result_dir="./tmp/"):
                     fn=enhance_with_gpu,  # Use the GPU-decorated function
                     cache_examples=False,
                 )
+
             end_btn.click(
-                inputs=[input_video, input_text],
+                inputs=[input_video, input_text, upscale, max_chunk_len, chunk_size],
                 outputs=[output_video],
                 fn=enhance_with_gpu,  # Use the GPU-decorated function
             )
 
     return star_iface
 
-
 if __name__ == "__main__":
     result_dir = os.path.join("./", "results")
     star_iface = star_demo(result_dir)

video_super_resolution/scripts/inference_sr.py

@@ -26,6 +26,7 @@ class STAR_sr():
                  upscale=4,
                  max_chunk_len=32,
                  variant_info=None,
+                 chunk_size=3,
                  ):
         self.model_path=model_path
         logger.info('checkpoint_path: {}'.format(self.model_path))
@@ -36,6 +37,7 @@ class STAR_sr():
 
         model_cfg = EasyDict(__name__='model_cfg')
         model_cfg.model_path = self.model_path
+        model_cfg.chunk_size = chunk_size
         self.model = VideoToVideo_sr(model_cfg)
 
         steps = 15 if solver_mode == 'fast' else steps

video_to_video/video_to_video_model.py

@@ -88,6 +88,8 @@ class VideoToVideo_sr():
 
         negative_y = text_encoder(self.negative_prompt).detach()
         self.negative_y = negative_y
+        
+        self.chunk_size = opt.chunk_size
 
 
     def test(self, input: Dict[str, Any], total_noise_levels=1000, \
@@ -142,7 +144,7 @@ class VideoToVideo_sr():
             torch.cuda.empty_cache()
 
             logger.info(f'sampling, finished.')
-            vid_tensor_gen = self.vae_decode_chunk(gen_vid, chunk_size=3)
+            vid_tensor_gen = self.vae_decode_chunk(gen_vid, chunk_size=self.chunk_size)
 
             logger.info(f'temporal vae decoding, finished.')