Update app.py

app.py

@@ -508,71 +508,93 @@ def generate_video():
                     ref_pose_video.save(ref_pose_video_path)
                     print('ref_pose_video_path',ref_pose_video_path)
                     
+    
+        # Example of using the class with some hypothetical paths
+        args = AnimationConfig(driven_audio_path=driven_audio_path, source_image_path=source_image_path, result_folder=result_folder, pose_style=pose_style, expression_scale=expression_scale,enhancer=enhancer,still=still,preprocess=preprocess,ref_pose_video_path=ref_pose_video_path, image_hardcoded=image_hardcoded)
+            
+        if torch.cuda.is_available() and not args.cpu:
+            args.device = "cuda"
+        else:
+            args.device = "cpu"
+            
+        generation_thread = threading.Thread(target=main, args=(args,))
+        app.config['generation_thread'] = generation_thread
+        generation_thread.start()
+        # response_data = {"message": "Video generation started",
+        #                 "process_id": generation_thread.ident}
+    
+        # return jsonify(response_data)
+    
+        while generation_thread.is_alive():
+            if app.config.get('temp_response'):
+                final_response = app.config['temp_response']
+                response_data = {
+                    "base64_video": final_response,
+                    "text_prompt": app.config.get('text_prompt'),
+                    "duration": app.config.get('final_video_duration'),
+                    "status": "completed"
+                }
+                final_video_path = app.config['final_video_path']
+                print('final_video_path', final_video_path)
+    
+                if final_video_path and os.path.exists(final_video_path):
+                    os.remove(final_video_path)
+                    print("Deleted video file:", final_video_path)
+    
+                preprocess_dir = os.path.join("/tmp", "preprocess_data")
+                custom_cleanup(TEMP_DIR.name, preprocess_dir)
+    
+                end_time = time.time()
+                total_time = round(end_time - start_time, 2)
+                print("Total time taken for execution:", total_time, " seconds")
+                response_data["time_taken"] = total_time
+    
+                return jsonify(response_data)
+
     except Exception as e:
         app.logger.error(f"An error occurred: {e}")
         return "An error occurred", 500
-    
-    # Example of using the class with some hypothetical paths
-    args = AnimationConfig(driven_audio_path=driven_audio_path, source_image_path=source_image_path, result_folder=result_folder, pose_style=pose_style, expression_scale=expression_scale,enhancer=enhancer,still=still,preprocess=preprocess,ref_pose_video_path=ref_pose_video_path, image_hardcoded=image_hardcoded)
-        
-    if torch.cuda.is_available() and not args.cpu:
-        args.device = "cuda"
-    else:
-        args.device = "cpu"
-        
-    generation_thread = threading.Thread(target=main, args=(args,))
-    app.config['generation_thread'] = generation_thread
-    generation_thread.start()
-    response_data = {"message": "Video generation started",
-                    "process_id": generation_thread.ident}
 
-    return jsonify(response_data)
-    # base64_video = main(args)
-    # return jsonify({"base64_video": base64_video})
-
-    #else:
-    #    return 'Unsupported HTTP method', 405
-
-@app.route("/status", methods=["GET"])
-def check_generation_status():
-    global TEMP_DIR
-    global start_time
-    response = {"base64_video": "","text_prompt":"", "status": ""}
-    process_id = request.args.get('process_id', None)
-
-    # process_id is required to check the status for that specific process
-    if process_id:
-        generation_thread = app.config.get('generation_thread')
-        if generation_thread and generation_thread.ident == int(process_id) and generation_thread.is_alive():
-            return jsonify({"status": "in_progress"}), 200
-        elif app.config.get('temp_response'):
-            # app.config['temp_response']['status'] = 'completed'
-            final_response = app.config['temp_response']
-            response["base64_video"] = final_response
-            response["text_prompt"] = app.config.get('text_prompt')
-            response["duration"] = app.config.get('final_video_duration')
-            response["status"] = "completed"
-
-            final_video_path = app.config['final_video_path']
-            print('final_video_path',final_video_path)
-
-
-            if final_video_path and os.path.exists(final_video_path):
-                os.remove(final_video_path)
-                print("Deleted video file:", final_video_path)
-
-            # TEMP_DIR.cleanup()
-            preprocess_dir = os.path.join("/tmp", "preprocess_data")
-            custom_cleanup(TEMP_DIR.name, preprocess_dir)
-
-            print("Temporary files cleaned up, but preprocess_data is retained.")
+# @app.route("/status", methods=["GET"])
+# def check_generation_status():
+#     global TEMP_DIR
+#     global start_time
+#     response = {"base64_video": "","text_prompt":"", "status": ""}
+#     process_id = request.args.get('process_id', None)
+
+#     # process_id is required to check the status for that specific process
+#     if process_id:
+#         generation_thread = app.config.get('generation_thread')
+#         if generation_thread and generation_thread.ident == int(process_id) and generation_thread.is_alive():
+#             return jsonify({"status": "in_progress"}), 200
+#         elif app.config.get('temp_response'):
+#             # app.config['temp_response']['status'] = 'completed'
+#             final_response = app.config['temp_response']
+#             response["base64_video"] = final_response
+#             response["text_prompt"] = app.config.get('text_prompt')
+#             response["duration"] = app.config.get('final_video_duration')
+#             response["status"] = "completed"
+
+#             final_video_path = app.config['final_video_path']
+#             print('final_video_path',final_video_path)
+
+
+#             if final_video_path and os.path.exists(final_video_path):
+#                 os.remove(final_video_path)
+#                 print("Deleted video file:", final_video_path)
+
+#             # TEMP_DIR.cleanup()
+#             preprocess_dir = os.path.join("/tmp", "preprocess_data")
+#             custom_cleanup(TEMP_DIR.name, preprocess_dir)
+
+#             print("Temporary files cleaned up, but preprocess_data is retained.")
             
-            end_time = time.time()
-            total_time = round(end_time - start_time, 2)
-            print("Total time taken for execution:", total_time, " seconds")
-            response["time_taken"] = total_time
-            return jsonify(response)
-    return jsonify({"error":"No process id provided"})
+#             end_time = time.time()
+#             total_time = round(end_time - start_time, 2)
+#             print("Total time taken for execution:", total_time, " seconds")
+#             response["time_taken"] = total_time
+#             return jsonify(response)
+#     return jsonify({"error":"No process id provided"})
 
 @app.route("/health", methods=["GET"])
 def health_status():