Merge pull request #65 from bghira/main

Logging reduction and S3 connection pool size fix
bghira · Aug 28, 2023 · 7d04f3f · 7d04f3f
2 parents 36fe422 + cb9d62f
commit 7d04f3f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/helpers/data_backend/aws.py b/helpers/data_backend/aws.py
@@ -3,6 +3,7 @@
 from torch import Tensor
 from pathlib import PosixPath
 import concurrent.futures
+from botocore.config import Config
 from helpers.data_backend.base import BaseDataBackend
 
 loggers_to_silence = [
@@ -49,10 +50,12 @@ def __init__(
             extra_args = {
                 "endpoint_url": endpoint_url,
             }
+        s3_config = Config(max_pool_connections=100)
         self.client = boto3.client(
             "s3",
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
+            config=s3_config,
             **extra_args,
         )
 

diff --git a/helpers/log_format.py b/helpers/log_format.py
@@ -32,5 +32,8 @@ def format(self, record):
 if not logger.handlers:
     logger.addHandler(new_handler)
 
+forward_logger = logging.getLogger('diffusers.models.unet_2d_condition')
+forward_logger.setLevel(logging.WARNING)
+
 pil_logger = logging.getLogger("PIL")
 pil_logger.setLevel(logging.INFO)