From 8ff9a648882d97f9dedfb9c381fb0893757adda2 Mon Sep 17 00:00:00 2001
From: tzwm <git@tzwm.me>
Date: Tue, 21 May 2024 16:10:47 +0800
Subject: [PATCH] add retry to http_client

---
 nodes/dify_text_generator.py |  4 ++--
 nodes/load_image_by_url.py   |  4 ++--
 nodes/upload_to_remote.py    |  5 ++---
 nodes/xyz_plot.py            |  5 ++---
 requirements.txt             |  2 ++
 routes/downloads.py          |  5 ++---
 utils.py                     | 11 +++++++++++
 7 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/nodes/dify_text_generator.py b/nodes/dify_text_generator.py
index 738c881..04229b7 100644
--- a/nodes/dify_text_generator.py
+++ b/nodes/dify_text_generator.py
@@ -1,5 +1,5 @@
-import requests
 import json
+from ..utils import http_client
 
 class DifyTextGenerator:
     CATEGORY = "Browser"
@@ -39,7 +39,7 @@ def run(self, dify_api_endpoint, api_key, query, inputs_json_str=None):
             # something weird, I have to add '{' and '}' manually
             data["inputs"] = json.loads("{" + inputs_json_str + "}")
 
-        r = requests.post(
+        r = http_client().post(
             dify_api_endpoint,
             headers=header,
             data=json.dumps(data),
diff --git a/nodes/load_image_by_url.py b/nodes/load_image_by_url.py
index a6e8f73..035dfee 100644
--- a/nodes/load_image_by_url.py
+++ b/nodes/load_image_by_url.py
@@ -1,9 +1,9 @@
 import hashlib
-import requests
 import os
 from PIL import Image, ImageSequence, ImageOps
 import numpy as np
 import torch
+from ..utils import http_client
 
 
 import folder_paths
@@ -36,7 +36,7 @@ def filename(self):
 
     def download_by_url(self):
         input_dir = folder_paths.get_input_directory()
-        res = requests.get(self.url)
+        res = http_client().get(self.url)
         if res.status_code == 200:
             download_path = os.path.join(input_dir, self.filename())
             with open(download_path, 'wb') as file:
diff --git a/nodes/upload_to_remote.py b/nodes/upload_to_remote.py
index cd8c9d9..dc64ebd 100644
--- a/nodes/upload_to_remote.py
+++ b/nodes/upload_to_remote.py
@@ -1,4 +1,3 @@
-import requests
 import threading
 import asyncio
 import json
@@ -7,7 +6,7 @@
 import numpy as np
 import io
 import base64
-from ..utils import log
+from ..utils import http_client, log
 
 class UploadToRemote:
     CATEGORY = "Browser"
@@ -89,7 +88,7 @@ async def callback(images, extra, remote_url, extension='jpeg', quality=85, embe
             }
             data = json.dumps(data).encode('utf-8')
             log(f"uploading {track_id} to {remote_url}")
-            res = requests.post(remote_url, data=data, headers=headers)
+            res = http_client().post(remote_url, data=data, headers=headers)
             log(f"uploaded {track_id}: {res.status_code} {res.text}")
             # TODO: check the response
 
diff --git a/nodes/xyz_plot.py b/nodes/xyz_plot.py
index 8282ba5..ae1017d 100644
--- a/nodes/xyz_plot.py
+++ b/nodes/xyz_plot.py
@@ -1,6 +1,5 @@
 import shutil
 import time
-import requests
 import json
 from PIL import Image
 import numpy as np
@@ -9,7 +8,7 @@
 
 import folder_paths
 
-from ..utils import SERVER_BASE_URL
+from ..utils import SERVER_BASE_URL, http_client
 
 class XyzPlot:
     CATEGORY = "Browser"
@@ -107,7 +106,7 @@ def queue_new_prompt(prompt):
 
             # for some special network environments like AutoDL
             proxies = {"http": "", "https": ""}
-            return requests.post(SERVER_BASE_URL + '/prompt', data=data, proxies=proxies)
+            return http_client().post(SERVER_BASE_URL + '/prompt', data=data, proxies=proxies)
 
 
         batch_size = len(images)
diff --git a/requirements.txt b/requirements.txt
index bfb7764..e17c135 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,5 @@
 tqdm
 pandas
 numpy
+requests
+urllib3
diff --git a/routes/downloads.py b/routes/downloads.py
index d98fcc0..43c770f 100644
--- a/routes/downloads.py
+++ b/routes/downloads.py
@@ -1,7 +1,6 @@
 from os import path
 import os
 import shutil
-import requests
 import time
 import asyncio
 import json
@@ -11,7 +10,7 @@
 
 import folder_paths
 
-from ..utils import download_logs_path, log
+from ..utils import download_logs_path, log, http_client
 
 def parse_options_header(content_disposition):
     param, options = '', {}
@@ -53,7 +52,7 @@ async def download_by_requests(uuid:str, download_url:str, save_in:str, filename
 
     HEADERS = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36"}
 
-    with requests.get(download_url, headers=HEADERS, stream=True) as resp:
+    with http_client().get(download_url, headers=HEADERS, stream=True) as resp:
         MISSING_FILENAME = f"unkwown_{uuid}"
         # get file name
         if filename == "":
diff --git a/utils.py b/utils.py
index 70a32fb..d75366e 100644
--- a/utils.py
+++ b/utils.py
@@ -4,6 +4,8 @@
 import subprocess
 import time
 from typing import TypedDict, List
+import requests
+from requests.adapters import HTTPAdapter, Retry
 
 import folder_paths
 from comfy.cli_args import args
@@ -24,6 +26,15 @@
 
 git_remote_name = 'origin'
 
+def http_client():
+    adapter = HTTPAdapter(max_retries=Retry(3, backoff_factor=0.1))
+    http = requests.session()
+    http.mount('http://', adapter)
+    http.mount('https://', adapter)
+
+    return http
+
+
 @functools.cache
 def get_config():
     return {