From 1ef30eeddf1186068abb00586d47097778542255 Mon Sep 17 00:00:00 2001
From: gaclove <peng.gaoc@gmail.com>
Date: Tue, 15 Jul 2025 00:37:52 +0800
Subject: [PATCH] feat: enhance LightX2VInference with environment variable
 configuration, improved error handling, and temporary file management for
 better robustness in video generation

---
 lightx2v_nodes/universal_bridge.py | 48 ++++++++++++++++++++++++------
 1 file changed, 39 insertions(+), 9 deletions(-)

diff --git a/lightx2v_nodes/universal_bridge.py b/lightx2v_nodes/universal_bridge.py
index cb0dcd4..04bfe97 100644
--- a/lightx2v_nodes/universal_bridge.py
+++ b/lightx2v_nodes/universal_bridge.py
@@ -757,12 +757,22 @@ class LightX2VInference:
         config.negative_prompt = negative_prompt
         config.mode = "infer"
 
+        # 设置环境变量（从run_wan_i2v.sh中提取）
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        if "DTYPE" not in os.environ:
+            os.environ["DTYPE"] = "BF16"  # 可通过配置更改
+        if "ENABLE_GRAPH_MODE" not in os.environ:
+            os.environ["ENABLE_GRAPH_MODE"] = "false"
+        if "ENABLE_PROFILING_DEBUG" not in os.environ:
+            os.environ["ENABLE_PROFILING_DEBUG"] = "true"
+
         # 处理输出路径
         if not hasattr(config, "save_video_path") or config.save_video_path is None:
             config.save_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
 
         # 临时文件列表，用于清理
         temp_files = []
+        temp_files.append(config.save_video_path)  # 添加输出视频到临时文件列表
 
         try:
             # 处理i2v任务的图像输入
@@ -773,6 +783,8 @@ class LightX2VInference:
                     pil_image.save(tmp.name)
                     config.image_path = tmp.name
                     temp_files.append(tmp.name)
+            elif config.task == "i2v" and image is None:
+                raise ValueError("i2v任务需要输入图像")
 
             # 处理音频输入
             if audio is not None and "audio" in config.model_cls:
@@ -790,34 +802,52 @@ class LightX2VInference:
             # 获取或创建runner
             runner = self.bridge.get_runner(config)
 
-            assert runner is not None, "Runner is None"
+            if runner is None:
+                raise RuntimeError("Failed to initialize runner")
 
             # 运行生成 - run_pipeline 是异步方法
             if asyncio.iscoroutinefunction(runner.run_pipeline):
                 # 在同步环境中运行异步方法
-                loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(loop)
+                try:
+                    # 使用现有的事件循环（如果存在）或创建新的
+                    loop = asyncio.get_event_loop()
+                    if loop.is_closed():
+                        loop = asyncio.new_event_loop()
+                        asyncio.set_event_loop(loop)
+                except RuntimeError:
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+
                 try:
                     result = loop.run_until_complete(runner.run_pipeline())
-                finally:
-                    loop.close()
+                except Exception as e:
+                    print(f"Error during pipeline execution: {e}")
+                    raise
             else:
                 result = runner.run_pipeline()
 
             # 检查结果是否已经保存到文件
             if os.path.exists(config.save_video_path):
+                # 不要删除输出文件，从临时文件列表中移除
+                temp_files.remove(config.save_video_path)
                 return (self.converter.video_to_latent(config.save_video_path),)
-            elif hasattr(result, "save_video_path") and os.path.exists(result.save_video_path):
+            elif result is not None and hasattr(result, "save_video_path") and os.path.exists(result.save_video_path):
                 return (self.converter.video_to_latent(result.save_video_path),)
             else:
-                # 假设张量输出
-                return (self.converter.tensor_to_latent(result),)
+                # 如果没有保存视频，抛出错误
+                raise RuntimeError("Video generation failed: no output video found")
 
+        except Exception as e:
+            print(f"Error in LightX2V generation: {e}")
+            raise
         finally:
             # 清理临时文件
             for temp_file in temp_files:
                 if os.path.exists(temp_file):
-                    os.unlink(temp_file)
+                    try:
+                        os.unlink(temp_file)
+                    except Exception:
+                        pass
 
 
 class LightX2VDistillConfig: