From 7e468bf7415831bf8307406af1f42facae3930fd Mon Sep 17 00:00:00 2001
From: zhonggang <54022200+zhonggang666@users.noreply.github.com>
Date: Thu, 9 Jul 2026 17:27:44 +0800
Subject: [PATCH] fix: add generated media download timeouts

---
 .../impl/base_image_generate_node.py                       | 7 ++++++-
 .../impl/base_image_to_video_node.py                       | 7 ++++++-
 .../impl/base_text_to_video_node.py                        | 7 ++++++-
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/apps/application/flow/step_node/image_generate_step_node/impl/base_image_generate_node.py b/apps/application/flow/step_node/image_generate_step_node/impl/base_image_generate_node.py
index 281122364be..511237d11b7 100644
--- a/apps/application/flow/step_node/image_generate_step_node/impl/base_image_generate_node.py
+++ b/apps/application/flow/step_node/image_generate_step_node/impl/base_image_generate_node.py
@@ -14,6 +14,9 @@
 from oss.serializers.file import FileSerializer
 
 
+GENERATED_MEDIA_DOWNLOAD_TIMEOUT = 30
+
+
 class BaseImageGenerateNode(IImageGenerateNode):
     def save_context(self, details, workflow_manage):
         self.context['answer'] = details.get('answer')
@@ -59,7 +62,9 @@ def execute(self, model_id, prompt, negative_prompt, dialogue_number, dialogue_t
             if isinstance(image_url, str):
                 if image_url.startswith('http'):
                     # HTTP URL 情况
-                    image_url = requests.get(image_url).content
+                    response = requests.get(image_url, timeout=GENERATED_MEDIA_DOWNLOAD_TIMEOUT)
+                    response.raise_for_status()
+                    image_url = response.content
                 elif image_url.startswith('data:image'):
                     # Data URL 格式 (data:image/png;base64,...)
                     import base64
diff --git a/apps/application/flow/step_node/image_to_video_step_node/impl/base_image_to_video_node.py b/apps/application/flow/step_node/image_to_video_step_node/impl/base_image_to_video_node.py
index b59147323b9..7034fe62d90 100644
--- a/apps/application/flow/step_node/image_to_video_step_node/impl/base_image_to_video_node.py
+++ b/apps/application/flow/step_node/image_to_video_step_node/impl/base_image_to_video_node.py
@@ -16,6 +16,9 @@
 from models_provider.tools import get_model_instance_by_model_workspace_id
 
 
+GENERATED_MEDIA_DOWNLOAD_TIMEOUT = 30
+
+
 class BaseImageToVideoNode(IImageToVideoNode):
     def save_context(self, details, workflow_manage):
         self.context['answer'] = details.get('answer')
@@ -64,7 +67,9 @@ def execute(self, model_id, prompt, negative_prompt, dialogue_number, dialogue_t
             return NodeResult({'answer': gettext('Failed to generate video')}, {})
         file_name = 'generated_video.mp4'
         if isinstance(video_urls, str) and video_urls.startswith('http'):
-            video_urls = requests.get(video_urls).content
+            response = requests.get(video_urls, timeout=GENERATED_MEDIA_DOWNLOAD_TIMEOUT)
+            response.raise_for_status()
+            video_urls = response.content
         file = bytes_to_uploaded_file(video_urls, file_name)
         file_url = self.upload_file(file)
         video_label = f'<video src="{file_url}" controls style="max-width: 100%; width: 100%; height: auto; max-height: 60vh;"></video>'
diff --git a/apps/application/flow/step_node/text_to_video_step_node/impl/base_text_to_video_node.py b/apps/application/flow/step_node/text_to_video_step_node/impl/base_text_to_video_node.py
index 806aa0116db..b04178c6cbe 100644
--- a/apps/application/flow/step_node/text_to_video_step_node/impl/base_text_to_video_node.py
+++ b/apps/application/flow/step_node/text_to_video_step_node/impl/base_text_to_video_node.py
@@ -15,6 +15,9 @@
 from django.utils.translation import gettext_lazy as _, gettext
 
 
+GENERATED_MEDIA_DOWNLOAD_TIMEOUT = 30
+
+
 class BaseTextToVideoNode(ITextToVideoNode):
     def save_context(self, details, workflow_manage):
         self.context['answer'] = details.get('answer')
@@ -57,7 +60,9 @@ def execute(self, model_id, prompt, negative_prompt, dialogue_number, dialogue_t
             return NodeResult({'answer': gettext('Failed to generate video')}, {})
         file_name = 'generated_video.mp4'
         if isinstance(video_urls, str) and video_urls.startswith('http'):
-            video_urls = requests.get(video_urls).content
+            response = requests.get(video_urls, timeout=GENERATED_MEDIA_DOWNLOAD_TIMEOUT)
+            response.raise_for_status()
+            video_urls = response.content
         file = bytes_to_uploaded_file(video_urls, file_name)
         file_url = self.upload_file(file)
         video_label = f'<video src="{file_url}" controls style="max-width: 100%; width: 100%; height: auto;"></video>'