curl with s3 private

hydra_node/config.py

@@ -139,6 +139,11 @@ def init_config_model():
     config.model_name = os.environ['MODEL']
     logger.info(f"MODEL: {config.model_name}")
     
+    #S3 auth stuff
+    config.s3_access_key = os.environ['S3_ACCESS_KEY']
+    config.s3_secret_key = os.environ['S3_SECRET_KEY']
+    config.s3_bucket = os.environ['S3_BUCKET']
+
     # Resolve where we get our model and data from.
     config.model_path = os.getenv('MODEL_PATH', None)
     config.enable_ema = os.getenv('ENABLE_EMA', "1")

hydra_node/sanitize.py

@@ -135,6 +135,10 @@ def sanitize_stable_diffusion(request, config):
         try:
             from PIL import Image
             image = Image.open(BytesIO(request.image))
+            #check if image resolution is bigger than 1280*1280, if discard
+            if image.width * image.height > 1280*1280:
+                return False, "image resolution is too big"
+                
             image.verify()
 
         except Exception as e:

hydra_node/web.py

@@ -17,6 +17,16 @@ import json
 import os
 from urllib import request
 import tempfile
+import os
+import subprocess
+import sys
+import time
+import datetime
+import hashlib
+import hmac
+import base64
+import requests
+import argparse
 
 from collections import OrderedDict
 from typing import Tuple, Union, List, Iterator, Callable
@@ -50,17 +60,42 @@ def convert_bytes(num):
             return "%3.1f %s" % (num, x)
         num /= step_unit
 
-class CURLStreamFile(object):
+def get_s3_secret_headers(endpoint, access_key, secret_key, bucket, s3_file):
+    resource=f"/{bucket}/{s3_file}"
+    content_type="binary/octet-stream"
+    date_value=datetime.datetime.utcnow().strftime('%a, %d %b %Y %H:%M:%S GMT')
+    string_to_sign="GET\n\n{}\n{}\n{}".format(content_type, date_value, resource)
+    #convert bash the last line into python
+    signature_hash=base64.b64encode(hmac.new(secret_key.encode(), string_to_sign.encode(), hashlib.sha1).digest()).decode()
+    headers = {
+        "Host": endpoint,
+        "Date": date_value,
+        "Content-Type": content_type,
+        "Authorization": "AWS {}:{}".format(access_key, signature_hash)
+    }
+    return headers
+
+class CURLS3StreamFile(object):
     """
     CURLStreamFile implements a file-like object around an HTTP download, the
     intention being to not buffer more than we have to.
     """
 
-    def __init__(self, uri: str) -> None:
-        # NOTE: `256mb` buffer on the python IO object.
-        self._curl = subprocess.Popen(['/usr/bin/curl',
-                                       '--header', 'Accept-Encoding: identity',
-                                       '-s', uri],
+    def __init__(self, url, headers=None) -> None:
+        subprocess_list = [
+            '/usr/bin/curl'
+        ]
+        for k, v in headers.items():
+            subprocess_list.append('-H')
+            subprocess_list.append(f"{k}: {v}")
+
+        subprocess_list.append(url)
+        subprocess_list.extend(['-H', 'Accept-Encoding: identity'])
+        subprocess_list.append('-s')
+        #subprocess_list.append('-o test124.yaml')
+        print(subprocess_list)
+
+        self._curl = subprocess.Popen(subprocess_list,
                                       stdout=subprocess.PIPE,
                                       bufsize=256 * 1024 * 1024)
         # Read our max-fd-size, fall back to 1mb if invalid.