Update app.py

app.py

@@ -3,41 +3,46 @@ import subprocess
 import os
 import requests
 
-def process_video(audio_file, video_file):
-    print(gradio.__version__)
-    # Unpack the audio and video file paths
-    audio_path = audio_file[1] if isinstance(audio_file, tuple) else audio_file
-    video_path = video_file if isinstance(video_file, str) else video_file.name
-    out_path = "output_video.mp4"
-    
-    # Define command flags
-    sample_mode = "cross"  # or "reconstruction"
-    generate_from_filelist = 0
-    model_path = "checkpoints/checkpoint.pt"
-    pads = "0,0,0,0"
-    
-    if sample_mode == "reconstruction":
-        sample_input_flags = "--sampling_input_type=first_frame --sampling_ref_type=first_frame"
-    elif sample_mode == "cross":
-        sample_input_flags = "--sampling_input_type=gt --sampling_ref_type=gt"
-    else:
-        return "Error: sample_mode can only be \"cross\" or \"reconstruction\""
-    
-    MODEL_FLAGS = "--attention_resolutions 32,16,8 --class_cond False --learn_sigma True --num_channels 128 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm False"
-    DIFFUSION_FLAGS = "--predict_xstart False --diffusion_steps 1000 --noise_schedule linear --rescale_timesteps False"
-    SAMPLE_FLAGS = f"--sampling_seed=7 {sample_input_flags} --timestep_respacing ddim25 --use_ddim True --model_path={model_path}"
-    DATA_FLAGS = "--nframes 5 --nrefer 1 --image_size 128 --sampling_batch_size=32"
-    TFG_FLAGS = "--face_hide_percentage 0.5 --use_ref=True --use_audio=True --audio_as_style=True"
-    GEN_FLAGS = f"--generate_from_filelist {generate_from_filelist} --video_path={video_path} --audio_path={audio_path} --out_path={out_path} --save_orig=False --face_det_batch_size 16 --pads {pads} --is_voxceleb2=False"
-    
-    # Combine all flags into one command
-    command = f"python your_model_script.py {MODEL_FLAGS} {DIFFUSION_FLAGS} {SAMPLE_FLAGS} {DATA_FLAGS} {TFG_FLAGS} {GEN_FLAGS}"
-    
+def process_video(audio_file, video_file, status):
     try:
+        # Unpack the audio and video file paths
+        audio_path = audio_file[1] if isinstance(audio_file, tuple) else audio_file
+        video_path = video_file if isinstance(video_file, str) else video_file.name
+        out_path = "output_video.mp4"
+        
+        # Define command flags
+        sample_mode = "cross"  # or "reconstruction"
+        generate_from_filelist = 0
+        model_path = "checkpoints/checkpoint.pt"
+        pads = "0,0,0,0"
+        
+        if sample_mode == "reconstruction":
+            sample_input_flags = "--sampling_input_type=first_frame --sampling_ref_type=first_frame"
+        elif sample_mode == "cross":
+            sample_input_flags = "--sampling_input_type=gt --sampling_ref_type=gt"
+        else:
+            return None, "Error: sample_mode can only be \"cross\" or \"reconstruction\""
+        
+        MODEL_FLAGS = "--attention_resolutions 32,16,8 --class_cond False --learn_sigma True --num_channels 128 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm False"
+        DIFFUSION_FLAGS = "--predict_xstart False --diffusion_steps 1000 --noise_schedule linear --rescale_timesteps False"
+        SAMPLE_FLAGS = f"--sampling_seed=7 {sample_input_flags} --timestep_respacing ddim25 --use_ddim True --model_path={model_path}"
+        DATA_FLAGS = "--nframes 5 --nrefer 1 --image_size 128 --sampling_batch_size=32"
+        TFG_FLAGS = "--face_hide_percentage 0.5 --use_ref=True --use_audio=True --audio_as_style=True"
+        GEN_FLAGS = f"--generate_from_filelist {generate_from_filelist} --video_path={video_path} --audio_path={audio_path} --out_path={out_path} --save_orig=False --face_det_batch_size 16 --pads {pads} --is_voxceleb2=False"
+        
+        # Combine all flags into one command
+        command = f"python your_model_script.py {MODEL_FLAGS} {DIFFUSION_FLAGS} {SAMPLE_FLAGS} {DATA_FLAGS} {TFG_FLAGS} {GEN_FLAGS}"
+        
+        # Execute command
         subprocess.run(command, shell=True, check=True)
-        return out_path
-    except subprocess.CalledProcessError as e:
-        return f"Error processing video: {e}"
+        
+        # If successful, return the output path and success message
+        return out_path, "Processing completed successfully!"
+        
+    except Exception as e:
+        # If there's an error, return None for the video and the error message
+        return None, f"Error during processing: {str(e)}"
+    
     finally:
         # Clean up output file if it exists
         if os.path.exists(out_path):
@@ -52,27 +57,21 @@ with gr.Blocks() as iface:
         audio_input = gr.Audio(label="Input Audio")
         video_input = gr.Video(label="Input Video")
     
-    with gr.Row():
-        process_button = gr.Button("Process Video")
-        status_msg = gr.Textbox(label="Status", interactive=False)
-    
+    status_msg = gr.Textbox(label="Status", interactive=False)
+    process_button = gr.Button("Process Video")
     video_output = gr.Video(label="Processed Video")
     
-    def process_with_status(audio, video):
-        try:
-            status_msg.update(value="Processing... Please wait.")
-            result = process_video(audio, video)
-            status_msg.update(value="Done!")
-            return [result, "Processing completed successfully!"]
-        except Exception as e:
-            error_msg = f"Error during processing: {str(e)}"
-            status_msg.update(value=error_msg)
-            return [None, error_msg]
-    
     process_button.click(
-        fn=process_with_status,
-        inputs=[audio_input, video_input],
-        outputs=[video_output, status_msg]
+        fn=process_video,
+        inputs=[
+            audio_input,
+            video_input,
+            status_msg
+        ],
+        outputs=[
+            video_output,
+            status_msg
+        ]
     )
 
 # Launch the interface