Update app.py

app.py

@@ -31,6 +31,7 @@ def download_model(repo_id, model_name):
     model_path = hf_hub_download(repo_id=repo_id, filename=model_name)
     return model_path
 
+import glob
 
 @spaces.GPU
 def run_inference(model_name, prompt_text):
@@ -51,7 +52,6 @@ def run_inference(model_name, prompt_text):
     prompt_file.write(prompt_text)
     prompt_file.close()
 
-    # Read and update the configuration file
     with open(config_path, 'r') as file:
         config_content = file.read()
     config_content = config_content.replace('prompt_path = "./assets/texts/t2v_samples.txt"', f'prompt_path = "{prompt_file.name}"')
@@ -60,24 +60,23 @@ def run_inference(model_name, prompt_text):
         temp_file.write(config_content)
         temp_config_path = temp_file.name
 
-
     cmd = [
         "torchrun", "--standalone", "--nproc_per_node", "1",
         "scripts/inference.py", temp_config_path,
         "--ckpt-path", ckpt_path
     ]
-    result = subprocess.run(cmd, capture_output=True, text=True)
+    subprocess.run(cmd)
+
+    # Find the latest generated file in the save directory
+    save_dir = "./output"  # Or the appropriate directory where files are saved
+    list_of_files = glob.glob(f'{save_dir}/*')  # You might need to adjust the pattern
+    latest_file = max(list_of_files, key=os.path.getctime)
+
     # Clean up the temporary files
     os.remove(temp_file.name)
     os.remove(prompt_file.name)
 
-    if result.returncode == 0:
-        # Assuming the output video is saved at a known location, for example "./output/video.mp4"
-        output_video_path = "./output/video.mp4"
-        return output_video_path
-    else:
-        print("Error occurred:", result.stderr)
-        return None  # You might want to handle errors differently
+    return latest_file
 
 def main():
     gr.Interface(