Add timeout to requests calls

llava/eval/run_llava.py

@@ -16,7 +16,7 @@
 
 def load_image(image_file):
     if image_file.startswith('http') or image_file.startswith('https'):
-        response = requests.get(image_file)
+        response = requests.get(image_file, timeout=60)
         image = Image.open(BytesIO(response.content)).convert('RGB')
     else:
         image = Image.open(image_file).convert('RGB')

llava/serve/cli.py

@@ -17,7 +17,7 @@
 
 def load_image(image_file):
     if image_file.startswith('http://') or image_file.startswith('https://'):
-        response = requests.get(image_file)
+        response = requests.get(image_file, timeout=60)
         image = Image.open(BytesIO(response.content)).convert('RGB')
     else:
         image = Image.open(image_file).convert('RGB')

llava/serve/gradio_web_server.py

@@ -36,9 +36,9 @@ def get_conv_log_filename():
 
 
 def get_model_list():
-    ret = requests.post(args.controller_url + "/refresh_all_workers")
+    ret = requests.post(args.controller_url + "/refresh_all_workers", timeout=60)
     assert ret.status_code == 200
-    ret = requests.post(args.controller_url + "/list_models")
+    ret = requests.post(args.controller_url + "/list_models", timeout=60)
     models = ret.json()["models"]
     models.sort(key=lambda x: priority.get(x, x))
     logger.info(f"Models: {models}")
@@ -198,7 +198,7 @@ def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request:
     # Query worker address
     controller_url = args.controller_url
     ret = requests.post(controller_url + "/get_worker_address",
-            json={"model": model_name})
+            json={"model": model_name}, timeout=60)
     worker_addr = ret.json()["address"]
     logger.info(f"model_name: {model_name}, worker_addr: {worker_addr}")
 

llava/serve/model_worker.py

@@ -81,7 +81,7 @@ def register_to_controller(self):
             "check_heart_beat": True,
             "worker_status": self.get_status()
         }
-        r = requests.post(url, json=data)
+        r = requests.post(url, json=data, timeout=60)
         assert r.status_code == 200
 
     def send_heart_beat(self):

llava/serve/register_worker.py

@@ -22,5 +22,5 @@
         "check_heart_beat": args.check_heart_beat,
         "worker_status": None,
     }
-    r = requests.post(url, json=data)
+    r = requests.post(url, json=data, timeout=60)
     assert r.status_code == 200

llava/serve/test_message.py

@@ -11,14 +11,14 @@ def main():
         worker_addr = args.worker_address
     else:
         controller_addr = args.controller_address
-        ret = requests.post(controller_addr + "/refresh_all_workers")
-        ret = requests.post(controller_addr + "/list_models")
+        ret = requests.post(controller_addr + "/refresh_all_workers", timeout=60)
+        ret = requests.post(controller_addr + "/list_models", timeout=60)
         models = ret.json()["models"]
         models.sort()
         print(f"Models: {models}")
 
         ret = requests.post(controller_addr + "/get_worker_address",
-            json={"model": args.model_name})
+            json={"model": args.model_name}, timeout=60)
         worker_addr = ret.json()["address"]
         print(f"worker_addr: {worker_addr}")
 
@@ -38,7 +38,7 @@ def main():
         "stop": conv.sep,
     }
     response = requests.post(worker_addr + "/worker_generate_stream", headers=headers,
-            json=pload, stream=True)
+            json=pload, stream=True, timeout=60)
 
     print(prompt.replace(conv.sep, "\n"), end="")
     for chunk in response.iter_lines(chunk_size=8192, decode_unicode=False, delimiter=b"\0"):

sgm/modules/autoencoding/lpips/util.py

@@ -15,7 +15,7 @@
 
 def download(url, local_path, chunk_size=1024):
     os.makedirs(os.path.split(local_path)[0], exist_ok=True)
-    with requests.get(url, stream=True) as r:
+    with requests.get(url, stream=True, timeout=60) as r:
         total_size = int(r.headers.get("content-length", 0))
         with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
             with open(local_path, "wb") as f: