Update app.py

app.py

@@ -1,78 +1,65 @@
 import gradio as gr
-import subprocess
 import os
-import requests
+import subprocess
 
-def process_video(audio_file, video_file, status):
+# Replace with your model loading and processing function
+def process_audio_video(audio_file, video_file):
+    audio_path = "input_audio.wav"
+    video_path = "input_video.mp4"
+    out_path = "output_video.mp4"
+    
+    # Save uploaded files
+    audio_file.save(audio_path)
+    video_file.save(video_path)
+    
+    # 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}"
+    
     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 generate.py {MODEL_FLAGS} {DIFFUSION_FLAGS} {SAMPLE_FLAGS} {DATA_FLAGS} {TFG_FLAGS} {GEN_FLAGS}"
-        
-        # Execute command
         subprocess.run(command, shell=True, check=True)
-        
-        # 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)}"
-    
+        return out_path
+    except subprocess.CalledProcessError as e:
+        return f"Error processing video: {e}"
     finally:
-        # Clean up output file if it exists
+        # Clean up the files after processing
+        if os.path.exists(audio_path):
+            os.remove(audio_path)
+        if os.path.exists(video_path):
+            os.remove(video_path)
         if os.path.exists(out_path):
             os.remove(out_path)
 
-# Create a Gradio interface
-with gr.Blocks() as iface:
-    gr.Markdown("# Audio-Video Processing")
-    gr.Markdown("Upload an audio file and a video file to process the video based on the audio input.")
-    
-    with gr.Row():
-        audio_input = gr.Audio(label="Input Audio")
-        video_input = gr.Video(label="Input Video")
-    
-    status_msg = gr.Textbox(label="Status", interactive=False)
-    process_button = gr.Button("Process Video")
-    video_output = gr.Video(label="Processed Video")
-    
-    process_button.click(
-        fn=process_video,
-        inputs=[
-            audio_input,
-            video_input,
-            status_msg
-        ],
-        outputs=[
-            video_output,
-            status_msg
-        ]
-    )
+# Define the Gradio interface
+interface = gr.Interface(
+    fn=process_audio_video,
+    inputs=[
+        gr.Audio(label="Audio File"),
+        gr.Video(label="Video File"),
+    ],
+    outputs="video",
+    description="Process Audio and Video with your Model",
+    allow_flagging=False  # Disable flagging as output is a video
+)
 
-# Launch the interface
-iface.launch()
+# Launch the Gradio app
+interface.launch(share=True)