dropdown added

app.py

@@ -18,11 +18,15 @@ def process_video(
     guidance_scale,
     inference_steps,
     seed,
+    checkpoint_file,
 ):
     # Create the temp directory if it doesn't exist
     output_dir = Path("./temp")
     output_dir.mkdir(parents=True, exist_ok=True)
 
+    # Use selected checkpoint or fall back to default
+    checkpoint_path = Path("checkpoints/unetFiles") / checkpoint_file if checkpoint_file else CHECKPOINT_PATH
+
     # Convert paths to absolute Path objects and normalize them
     video_file_path = Path(video_path)
     video_path = video_file_path.absolute().as_posix()
@@ -44,7 +48,7 @@ def process_video(
     )
 
     # Parse the arguments
-    args = create_args(video_path, audio_path, output_path, guidance_scale, seed)
+    args = create_args(video_path, audio_path, output_path, guidance_scale, seed, checkpoint_path)
 
     try:
         result = main(
@@ -59,7 +63,7 @@ def process_video(
 
 
 def create_args(
-    video_path: str, audio_path: str, output_path: str, guidance_scale: float, seed: int
+    video_path: str, audio_path: str, output_path: str, guidance_scale: float, seed: int, checkpoint_path: Path
 ) -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--inference_ckpt_path", type=str, required=True)
@@ -72,7 +76,7 @@ def create_args(
     return parser.parse_args(
         [
             "--inference_ckpt_path",
-            CHECKPOINT_PATH.absolute().as_posix(),
+            checkpoint_path.absolute().as_posix(),
             "--video_path",
             video_path,
             "--audio_path",
@@ -86,6 +90,12 @@ def create_args(
         ]
     )
 
+# Add this function to get checkpoint files
+def get_checkpoint_files():
+    unet_files_dir = Path("unetFiles")
+    if not unet_files_dir.exists():
+        return []
+    return [f.name for f in unet_files_dir.glob("*.pt")]
 
 # Create Gradio interface
 with gr.Blocks(title="SoundImage") as demo:
@@ -99,6 +109,12 @@ with gr.Blocks(title="SoundImage") as demo:
 
     with gr.Row():
         with gr.Column():
+            # Add checkpoint selector dropdown
+            checkpoint_dropdown = gr.Dropdown(
+                choices=get_checkpoint_files(),
+                label="Select Checkpoint",
+                value=get_checkpoint_files()[0] if get_checkpoint_files() else None
+            )
             video_input = gr.Video(label="Input Video")
             audio_input = gr.Audio(label="Input Audio", type="filepath")
 
@@ -139,6 +155,7 @@ with gr.Blocks(title="SoundImage") as demo:
             guidance_scale,
             inference_steps,
             seed,
+            checkpoint_dropdown,
         ],
         outputs=video_output,
     )