Merge branch 'api' into api_airi_new_merge

Author: zhchaoxing

modules/api/api.py

@@ -38,11 +38,6 @@
 import uuid
 import os
 import json
-import boto3
-cache = dict()
-s3_client = boto3.client('s3')
-s3_resource= boto3.resource('s3')
-generated_images_s3uri = os.environ.get('generated_images_s3uri', None)
 
 def upscaler_to_index(name: str):
     try:
@@ -107,6 +102,35 @@ def encode_pil_to_base64(image):
 
     return base64.b64encode(bytes_data)
 
+def export_pil_to_bytes(image):
+    with io.BytesIO() as output_bytes:
+
+        if opts.samples_format.lower() == 'png':
+            use_metadata = False
+            metadata = PngImagePlugin.PngInfo()
+            for key, value in image.info.items():
+                if isinstance(key, str) and isinstance(value, str):
+                    metadata.add_text(key, value)
+                    use_metadata = True
+            image.save(output_bytes, format="PNG", pnginfo=(metadata if use_metadata else None), quality=opts.jpeg_quality)
+
+        elif opts.samples_format.lower() in ("jpg", "jpeg", "webp"):
+            parameters = image.info.get('parameters', None)
+            exif_bytes = piexif.dump({
+                "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") }
+            })
+            if opts.samples_format.lower() in ("jpg", "jpeg"):
+                image.save(output_bytes, format="JPEG", exif = exif_bytes, quality=opts.jpeg_quality)
+            else:
+                image.save(output_bytes, format="WEBP", exif = exif_bytes, quality=opts.jpeg_quality)
+
+        else:
+            raise HTTPException(status_code=500, detail="Invalid image format")
+
+        bytes_data = output_bytes.getvalue()
+
+    return bytes_data
+
 def api_middleware(app: FastAPI):
     rich_available = True
     try:
@@ -211,7 +235,7 @@ def __init__(self, app: FastAPI, queue_lock: Lock):
         self.add_api_route("/sdapi/v1/unload-checkpoint", self.unloadapi, methods=["POST"])
         self.add_api_route("/sdapi/v1/reload-checkpoint", self.reloadapi, methods=["POST"])
         self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=ScriptsList)
-        self.add_api_route("/invocations", self.invocations, methods=["POST"], response_model=Union[TextToImageResponse, ImageToImageResponse, ExtrasSingleImageResponse, ExtrasBatchImagesResponse,MemoryResponse,List[SDModelItem],List[UpscalerItem],OptionsModel,List[SamplerItem],FlagsModel,ProgressResponse])
+        self.add_api_route("/invocations", self.invocations, methods=["POST"], response_model=Union[TextToImageResponse, ImageToImageResponse, ExtrasSingleImageResponse, ExtrasBatchImagesResponse, InvocationsErrorResponse, InterrogateResponse, MemoryResponse, List[SDModelItem], List[UpscalerItem], OptionsModel, List[SamplerItem], FlagsModel, ProgressResponse])
         self.add_api_route("/ping", self.ping, methods=["GET"], response_model=PingResponse)
 
         self.default_script_arg_txt2img = []
@@ -715,28 +739,21 @@ def get_memory(self):
         return MemoryResponse(ram = ram, cuda = cuda)
 
     def post_invocations(self, b64images, quality):
-        if generated_images_s3uri:
-            bucket, key = self.get_bucket_and_key(generated_images_s3uri)
+        if shared.generated_images_s3uri:
+            bucket, key = shared.get_bucket_and_key(shared.generated_images_s3uri)
+            if key.endswith('/'):
+                key = key[ : -1]
             images = []
             for b64image in b64images:
-                image = decode_base64_to_image(b64image).convert('RGB')
-                output = io.BytesIO()
-
-                try:
-                    if not quality:
-                        quality = 95
-
-                    image.save(output, format='PNG', quality=quality)
-                except Exception:
-                    image.save(output, format='PNG', quality=95)
-
-                image_id = str(uuid.uuid4())
-                s3_client.put_object(
-                    Body=output.getvalue(),
+                bytes_data = export_pil_to_bytes(decode_base64_to_image(b64image))
+                image_id = datetime.datetime.now().strftime(f"%Y%m%d%H%M%S-{uuid.uuid4()}")
+                suffix = opts.samples_format.lower()
+                shared.s3_client.put_object(
+                    Body=bytes_data,
                     Bucket=bucket,
-                    Key=f'{key}/{image_id}.png'
+                    Key=f'{key}/{image_id}.{suffix}'
                 )
-                images.append(f's3://{bucket}/{key}/{image_id}.png')
+                images.append(f's3://{bucket}/{key}/{image_id}.{suffix}')
             return images
         else:
             return b64images
@@ -804,7 +821,7 @@ def invocations(self, req: InvocationsRequest):
             hypernetwork_s3uri = shared.cmd_opts.hypernetwork_s3uri
 
             if hypernetwork_s3uri !='':
-                self.download_s3files(hypernetwork_s3uri, shared.cmd_opts.hypernetwork_dir)
+                shared.s3_download(hypernetwork_s3uri, shared.cmd_opts.hypernetwork_dir)
                 shared.reload_hypernetworks()
 
             if req.options != None:
@@ -816,14 +833,14 @@ def invocations(self, req: InvocationsRequest):
 
             if req.task == 'text-to-image':
                 if embeddings_s3uri != '':
-                    self.download_s3files(embeddings_s3uri, shared.cmd_opts.embeddings_dir)
+                    shared.s3_download(embeddings_s3uri, shared.cmd_opts.embeddings_dir)
                     sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings()
                 response = self.text2imgapi(req.txt2img_payload)
                 response.images = self.post_invocations(response.images, quality)
                 return response
             elif req.task == 'image-to-image':
                 if embeddings_s3uri != '':
-                    self.download_s3files(embeddings_s3uri, shared.cmd_opts.embeddings_dir)
+                    shared.s3_download(embeddings_s3uri, shared.cmd_opts.embeddings_dir)
                     sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings()
                 response = self.img2imgapi(req.img2img_payload)
                 response.images = self.post_invocations(response.images, quality)
@@ -836,6 +853,10 @@ def invocations(self, req: InvocationsRequest):
                 response = self.extras_batch_images_api(req.extras_batch_payload)
                 response.images = self.post_invocations(response.images, quality)
                 return response
+            elif req.task == 'interrogate':
+                response = self.interrogateapi(req.interrogate_payload)
+                return response
+            
             elif req.task == 'get-progress':
                 response = self.progressapi(req.progress_payload)
                 print("____________getting progress result: ")
@@ -873,44 +894,8 @@ def invocations(self, req: InvocationsRequest):
             return InvocationsErrorResponse(error = str(e))
 
     def ping(self):
-        print('-------ping------')
         return {'status': 'Healthy'}
 
     def launch(self, server_name, port):
         self.app.include_router(self.router)
         uvicorn.run(self.app, host=server_name, port=port)
-
-    def get_bucket_and_key(self, s3uri):
-        pos = s3uri.find('/', 5)
-        bucket = s3uri[5 : pos]
-        key = s3uri[pos + 1 : ]
-        return bucket, key
-
-    def download_s3files(self, s3uri, path):
-        global cache
-
-        pos = s3uri.find('/', 5)
-        bucket = s3uri[5 : pos]
-        key = s3uri[pos + 1 : ]
-
-        s3_bucket = s3_resource.Bucket(bucket)
-        objs = list(s3_bucket.objects.filter(Prefix=key))
-
-        if os.path.isfile('cache'):
-            cache = json.load(open('cache', 'r'))
-
-        for obj in objs:
-            if obj.key == key:
-                continue
-            response = s3_client.head_object(
-                Bucket = bucket,
-                Key =  obj.key
-            )
-            obj_key = 's3://{0}/{1}'.format(bucket, obj.key)
-            if obj_key not in cache or cache[obj_key] != response['ETag']:
-                filename = obj.key[obj.key.rfind('/') + 1 : ]
-
-                s3_client.download_file(bucket, obj.key, os.path.join(path, filename))
-                cache[obj_key] = response['ETag']
-
-        json.dump(cache, open('cache', 'w'))

modules/api/models.py

@@ -302,6 +302,7 @@ class InvocationsRequest(BaseModel):
     img2img_payload: Optional[StableDiffusionImg2ImgProcessingAPI]
     extras_single_payload: Optional[ExtrasSingleImageRequest]
     extras_batch_payload: Optional[ExtrasBatchImagesRequest]
+    interrogate_payload: Optional[InterrogateRequest]
     progress_payload:Optional[ProgressRequest]
     post_options_payload:Optional[dict]
 

modules/cmd_args.py

@@ -109,3 +109,4 @@
 parser.add_argument('--dreambooth-config-id', default='', type=str, help='Dreambooth config ID')
 parser.add_argument('--embeddings-s3uri', default='', type=str, help='Embedding S3Uri')
 parser.add_argument('--hypernetwork-s3uri', default='', type=str, help='Hypernetwork S3Uri')
+parser.add_argument('--region-name', default='', type=str, help='Region name')

modules/script_callbacks.py

@@ -93,6 +93,7 @@ def __init__(self, imgs, cols, rows):
     callbacks_infotext_pasted=[],
     callbacks_script_unloaded=[],
     callbacks_before_ui=[],
+    callbacks_update_cn_models=[]
 )
 
 
@@ -224,6 +225,12 @@ def before_ui_callback():
         except Exception:
             report_exception(c, 'before_ui')
 
+def update_cn_models_callback():
+    for c in callback_map['callbacks_update_cn_models']:
+        try:
+            c.callback()
+        except Exception:
+            report_exception(c, 'callbacks_update_cn_models')
 
 def add_callback(callbacks, fun):
     stack = [x for x in inspect.stack() if x.filename != __file__]
@@ -247,6 +254,8 @@ def remove_callbacks_for_function(callback_func):
         for callback_to_remove in [cb for cb in callback_list if cb.callback == callback_func]:
             callback_list.remove(callback_to_remove)
 
+def on_update_cn_models(callback):
+    add_callback(callback_map['callbacks_update_cn_models'], callback)
 
 def on_app_started(callback):
     """register a function to be called when the webui started, the gradio `Block` component and

modules/shared.py

@@ -4,7 +4,7 @@
 import os
 import sys
 import time
-
+import threading
 from PIL import Image
 import gradio as gr
 import tqdm
@@ -15,9 +15,59 @@
 import modules.devices as devices
 from modules import localization, script_loading, errors, ui_components, shared_items, cmd_args
 from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir
+from botocore.exceptions import ClientError
+import glob
 
 demo = None
 
+models_s3_bucket = None
+s3_folder_sd = None
+s3_folder_cn = None
+s3_folder_lora = None
+syncLock = threading.Lock()
+tmp_models_dir = '/tmp/models'
+tmp_cache_dir = '/tmp/model_sync_cache'
+class ModelsRef:
+    def __init__(self):
+        self.models_ref = {}
+
+    def get_models_ref_dict(self):
+        return self.models_ref
+    
+    def add_models_ref(self, model_name):
+        if model_name in self.models_ref:
+            self.models_ref[model_name] += 1
+        else:
+            self.models_ref[model_name] = 0
+
+    def remove_model_ref(self,model_name):
+        if self.models_ref.get(model_name):
+            del self.models_ref[model_name]
+
+    def get_models_ref(self, model_name):
+        return self.models_ref.get(model_name)
+    
+    def get_least_ref_model(self):
+        sorted_models = sorted(self.models_ref.items(), key=lambda item: item[1])
+        if sorted_models:
+            least_ref_model, least_counter = sorted_models[0]
+            return least_ref_model,least_counter
+        else:
+            return None,None
+    
+    def pop_least_ref_model(self):
+        sorted_models = sorted(self.models_ref.items(), key=lambda item: item[1])
+        if sorted_models:
+            least_ref_model, least_counter = sorted_models[0]
+            del self.models_ref[least_ref_model]
+            return least_ref_model,least_counter
+        else:
+            return None,None
+        
+sd_models_Ref = ModelsRef()
+cn_models_Ref = ModelsRef()
+lora_models_Ref = ModelsRef()
+
 parser = cmd_args.parser
 
 script_loading.preload_extensions(extensions_dir, parser)
@@ -657,3 +707,99 @@ def html(filename):
             return file.read()
 
     return ""
+
+import boto3
+import requests
+
+cache = dict()
+region_name = boto3.session.Session().region_name if not cmd_opts.train else cmd_opts.region_name
+s3_client = boto3.client('s3', region_name=region_name)
+endpointUrl = s3_client.meta.endpoint_url
+s3_client = boto3.client('s3', endpoint_url=endpointUrl, region_name=region_name)
+s3_resource= boto3.resource('s3')
+generated_images_s3uri = os.environ.get('generated_images_s3uri', None)
+
+def get_bucket_and_key(s3uri):
+    pos = s3uri.find('/', 5)
+    bucket = s3uri[5 : pos]
+    key = s3uri[pos + 1 : ]
+    return bucket, key
+
+def s3_download(s3uri, path):
+    global cache
+
+    print('---path---', path)
+    os.system(f'ls -l {os.path.dirname(path)}')
+
+    pos = s3uri.find('/', 5)
+    bucket = s3uri[5 : pos]
+    key = s3uri[pos + 1 : ]
+
+    objects = []
+    paginator = s3_client.get_paginator('list_objects_v2')
+    page_iterator = paginator.paginate(Bucket=bucket, Prefix=key)
+    for page in page_iterator:
+        if 'Contents' in page:
+            for obj in page['Contents']:
+                objects.append(obj)
+        if 'NextContinuationToken' in page:
+            page_iterator = paginator.paginate(Bucket=bucket, Prefix=key,
+                                                ContinuationToken=page['NextContinuationToken'])
+
+    if os.path.isfile('cache'):
+        cache = json.load(open('cache', 'r'))
+
+    for obj in objects:
+        if obj['Size'] == 0:
+            continue
+        response = s3_client.head_object(
+            Bucket = bucket,
+            Key =  obj['Key']
+        )
+        obj_key = 's3://{0}/{1}'.format(bucket, obj['Key'])
+        if obj_key not in cache or cache[obj_key] != response['ETag']:
+            filename = obj['Key'][obj['Key'].rfind('/') + 1 : ]
+
+            s3_client.download_file(bucket, obj['Key'], os.path.join(path, filename))
+            cache[obj_key] = response['ETag']
+
+    json.dump(cache, open('cache', 'w'))
+
+def http_download(httpuri, path):
+    with requests.get(httpuri, stream=True) as r:
+        r.raise_for_status()
+        with open(path, 'wb') as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
+
+def upload_s3files(s3uri, file_path_with_pattern):
+    pos = s3uri.find('/', 5)
+    bucket = s3uri[5 : pos]
+    key = s3uri[pos + 1 : ]
+
+    try:
+        for file_path in glob.glob(file_path_with_pattern):
+            file_name = os.path.basename(file_path)
+            __s3file = f'{key}{file_name}'
+            print(file_path, __s3file)
+            s3_client.upload_file(file_path, bucket, __s3file)
+    except ClientError as e:
+        print(e)
+        return False
+    return True
+
+def upload_s3folder(s3uri, file_path):
+    pos = s3uri.find('/', 5)
+    bucket = s3uri[5 : pos]
+    key = s3uri[pos + 1 : ]
+
+    try:
+        for path, _, files in os.walk(file_path):
+            for file in files:
+                dest_path = path.replace(file_path,"")
+                __s3file = f'{key}{dest_path}/{file}'
+                __local_file = os.path.join(path, file)
+                print(__local_file, __s3file)
+                s3_client.upload_file(__local_file, bucket, __s3file)
+    except Exception as e:
+        print(e)