Add our changes

Author: zhchaoxing

.gitignore

@@ -13,7 +13,6 @@ __pycache__
 /gfpgan/weights/*.pth
 /ui-config.json
 /outputs
-/config.json
 /log
 /webui.settings.bat
 /embeddings
@@ -37,3 +36,9 @@ notification.mp3
 /node_modules
 /package-lock.json
 /.coverage*
+.DS_Store
+build.info
+**/.DS_Store
+models/Stable-diffusion/00v2.1_768-ema-pruned.safetensors
+models/Stable-diffusion/00v2.1_768-ema-pruned.yaml
+webui-user.bat

config.json

@@ -0,0 +1,26 @@
+{
+    "samples_format": "jpeg",
+    "control_net_model_config": "models/cldm_v21.yaml",
+    "control_net_max_models_num": 3,
+    "control_net_model_cache_size": 5,
+    "control_net_monocular_depth_optim": true,
+    "control_net_cfg_based_guidance": true,
+    "control_net_control_transfer": true,
+    "control_net_no_detectmap": true,
+    "control_net_detectmap_autosaving": false,
+    "control_net_only_midctrl_hires": true,
+    "control_net_allow_script_control": true,
+    "control_net_skip_img2img_processing": false,
+    "control_net_only_mid_control": false,
+    "control_net_sync_field_args": false,
+    "CLIP_stop_at_last_layers": 2,
+    "multiple_tqdm": false,
+    "auto_vae_precision": true,
+    "upcast_attn": true,
+    "enable_quantization": true,
+    "img_max_size_mp": 10.0,
+    "export_for_4chan": true,
+    "img_downscale_threshold": 6.0,
+    "target_side_length": 3600.0,
+    "randn_source": "CPU"
+}

modules/api/api.py

@@ -2,6 +2,7 @@
 import io
 import os
 import time
+import copy
 import datetime
 import uvicorn
 import ipaddress
@@ -129,6 +130,72 @@ def decode_base64_to_image(encoding):
         raise HTTPException(status_code=500, detail="Invalid encoded image") from e
 
 
+user_input_data = {}
+
+
+def set_img_exif_dict():
+    title = "AIRI"
+    date_taken = "2001:01:01 01:01:01"
+    global user_input_data
+    if "date_taken" in user_input_data:
+        date_taken = user_input_data['date_taken']
+    copyright = "© AIRI Lab. All Rights Reserved."
+    camera_maker = "AIRI Lab"
+    camera_model = "AIRI Model 1.0"
+    user_id = "AIRI tester"
+    if "user_id" in user_input_data:
+        user_id = user_input_data['user_id']
+    keywords = "Generated in AIRI platform. https://airilab.com"
+    description = "An image processed by the AIRI platform."
+    software = "AIRI Platform v1.0"
+    # imageid = "imageid?"
+    # imagenum = "imagenum?"
+    # seed = "seed?"
+    exif_dict = {
+        "0th": {
+            piexif.ImageIFD.ImageDescription: description.encode('utf-8'),
+            piexif.ImageIFD.Make: camera_maker.encode('utf-8'),
+            piexif.ImageIFD.Model: camera_model.encode('utf-8'),
+            piexif.ImageIFD.Copyright: copyright.encode('utf-8'),
+            piexif.ImageIFD.Artist: user_id.encode('utf-8'),
+            piexif.ImageIFD.ProcessingSoftware: software.encode('utf-8'),
+            piexif.ImageIFD.Software: software.encode('utf-8'),
+            piexif.ImageIFD.DateTime: date_taken.encode('utf-8'),
+            piexif.ImageIFD.HostComputer: software.encode('utf-8'),
+            # piexif.ImageIFD.ImageID: imageid.encode('utf-8'),
+            # piexif.ImageIFD.ImageNumber: imagenum.encode('utf-8'),
+            piexif.ImageIFD.ImageHistory: keywords.encode('utf-8'),
+            # piexif.ImageIFD.ImageResources: description.encode('utf-8'),
+            # piexif.ImageIFD.Noise: seed.encode('utf-8'),
+            piexif.ImageIFD.Predictor: camera_model.encode('utf-8'),
+            piexif.ImageIFD.OriginalRawFileData: keywords.encode('utf-8'),
+            # piexif.ImageIFD.OriginalRawFileName: imageid.encode('utf-8'),
+            piexif.ImageIFD.ProfileCopyright: copyright.encode('utf-8'),
+            piexif.ImageIFD.ProfileEmbedPolicy: software.encode('utf-8'),
+            piexif.ImageIFD.Rating: "5".encode('utf-8'),
+            piexif.ImageIFD.ProfileName: user_id.encode('utf-8'),
+            # piexif.ImageIFD.XPAuthor: user_id.encode('utf-8'),
+            # piexif.ImageIFD.XPTitle: title.encode('utf-8'),
+            # piexif.ImageIFD.XPKeywords: keywords.encode('utf-8'),
+            # piexif.ImageIFD.XPComment: description.encode('utf-8'),
+            # piexif.ImageIFD.XPSubject: copyright.encode('utf-8'),
+        },
+        "Exif": {
+            piexif.ExifIFD.DateTimeOriginal: date_taken.encode('utf-8'),
+            piexif.ExifIFD.CameraOwnerName: user_id.encode('utf-8'),
+            piexif.ExifIFD.DateTimeDigitized: date_taken.encode('utf-8'),
+            piexif.ExifIFD.DeviceSettingDescription: camera_model.encode('utf-8'),
+            piexif.ExifIFD.FileSource: keywords.encode('utf-8'),
+            # piexif.ExifIFD.ImageUniqueID: imageid.encode('utf-8'),
+            piexif.ExifIFD.LensMake: camera_maker.encode('utf-8'),
+            piexif.ExifIFD.LensModel: camera_model.encode('utf-8'),
+            piexif.ExifIFD.MakerNote: description.encode('utf-8'),
+            piexif.ExifIFD.UserComment: description.encode('utf-8'),
+        }
+    }
+    return exif_dict
+
+
 def encode_pil_to_base64(image):
     with io.BytesIO() as output_bytes:
 
@@ -144,10 +211,14 @@ def encode_pil_to_base64(image):
         elif opts.samples_format.lower() in ("jpg", "jpeg", "webp"):
             if image.mode == "RGBA":
                 image = image.convert("RGB")
-            parameters = image.info.get('parameters', None)
-            exif_bytes = piexif.dump({
-                "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") }
-            })
+            # parameters = image.info.get('parameters', None)
+            # exif_bytes = piexif.dump({
+            #     "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") }
+            # })
+
+            # Convert dict to bytes
+            exif_bytes = piexif.dump(set_img_exif_dict())
+
             if opts.samples_format.lower() in ("jpg", "jpeg"):
                 image.save(output_bytes, format="JPEG", exif = exif_bytes, quality=opts.jpeg_quality)
             else:
@@ -381,6 +452,12 @@ def init_script_args(self, request, default_script_args, selectable_scripts, sel
 
         # Now check for always on scripts
         if request.alwayson_scripts:
+            global user_input_data
+            user_input_data = {}
+            if "user_input" in request.alwayson_scripts:
+                user_input_data = request.alwayson_scripts["user_input"]
+                request.alwayson_scripts.pop("user_input")
+
             for alwayson_script_name in request.alwayson_scripts.keys():
                 alwayson_script = self.get_script(alwayson_script_name, script_runner)
                 if alwayson_script is None:
@@ -535,7 +612,7 @@ def pnginfoapi(self, req: models.PNGInfoRequest):
         if(not req.image.strip()):
             return models.PNGInfoResponse(info="")
 
-        image = decode_base64_to_image(req.image.strip())
+        image = decode_to_image(req.image.strip())
         if image is None:
             return models.PNGInfoResponse(info="")
 
@@ -580,7 +657,7 @@ def interrogateapi(self, interrogatereq: models.InterrogateRequest):
         if image_b64 is None:
             raise HTTPException(status_code=404, detail="Image not found")
 
-        img = decode_base64_to_image(image_b64)
+        img = decode_to_image(image_b64)
         img = img.convert('RGB')
 
         # Override object param
@@ -623,6 +700,9 @@ def get_config(self):
 
         return options
 
+    def get_all_config(self):
+        return shared.opts.data
+
     def set_config(self, req: Dict[str, Any]):
         checkpoint_name = req.get("sd_model_checkpoint", None)
         if checkpoint_name is not None and checkpoint_name not in checkpoint_aliases:
@@ -865,12 +945,53 @@ def post_invocations(self, b64images, quality):
         else:
             return b64images
 
+    def truncate_content(self, value, limit=1000):
+        if isinstance(value, str):  # Only truncate if the value is a string
+            if len(value) > limit:
+                return value[:limit] + '...'
+        return value
+
+    def req_logging(self, obj, indent=1):
+        if "__dict__" in dir(obj):  # if value is an object, dive into it
+            items = obj.__dict__.items()
+        elif isinstance(obj, dict):  # if value is a dictionary, get items
+            items = obj.items()
+        elif isinstance(obj, list):  # if value is a list, enumerate items
+            items = enumerate(obj)
+        else:  # if value is not an object or dict or list, just print it
+            print("  " * indent + f"{self.truncate_content(obj)}")
+            return
+
+        for attr, value in items:
+            if value is None or value == {} or value == []:
+                continue
+            if isinstance(value, (list, dict)) or "__dict__" in dir(value):
+                print("  " * indent + f"{attr}:")
+                self.req_logging(value, indent + 1)
+            else:
+                print("  " * indent + f"{attr}: {self.truncate_content(value)}")
+
     def invocations(self, req: models.InvocationsRequest):
         with self.invocations_lock:
-            print('-------invocation------')
-            print(req.task)
-
+            print("\n ----------------------------invocation--------------------------- ")
             try:
+                print("")
+                self.req_logging(req)
+            except Exception as e:
+                print("console Log ran into issue: ", e)
+            # print(f"log@{datetime.datetime.now().strftime(f'%Y%m%d%H%M%S')} req in invocations: {req}")
+            global user_input_data
+            user_input_data = {}
+            # if 'alwayson_scripts' in req:
+            #     if "user_input" in req.alwayson_scripts:
+            #         user_input_data = req.alwayson_scripts["user_input"]
+            #         req.alwayson_scripts.pop("user_input")
+            if req.user_input != None:
+                user_input_data = req.user_input
+                # print(f"log@{datetime.datetime.now().strftime(f'%Y%m%d%H%M%S')} user_input processed in invocations")
+                # req.pop('user_input', None)
+
+        try:
                 if req.vae != None:
                     shared.opts.data['sd_vae'] = req.vae
                     refresh_vae_list()
@@ -906,6 +1027,15 @@ def invocations(self, req: models.InvocationsRequest):
                         sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings()
                     response = self.text2imgapi(req.txt2img_payload)
                     response.images = self.post_invocations(response.images, quality)
+                    response.parameters.clear()
+                    oldinfo = json.loads(response.info)
+                    if "all_prompts" in oldinfo:
+                        oldinfo.pop("all_prompts", None)
+                    if "all_negative_prompts" in oldinfo:
+                        oldinfo.pop("all_negative_prompts", None)
+                    if "infotexts" in oldinfo:
+                        oldinfo.pop("infotexts", None)
+                    response.info = json.dumps(oldinfo)
                     return response
                 elif req.task == 'image-to-image':
                     response = requests.get('http://0.0.0.0:8080/controlnet/model_list', params={'update': True})
@@ -916,10 +1046,51 @@ def invocations(self, req: models.InvocationsRequest):
                         sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings()
                     response = self.img2imgapi(req.img2img_payload)
                     response.images = self.post_invocations(response.images, quality)
+                    response.parameters.clear()
+                    oldinfo = json.loads(response.info)
+                    if "all_prompts" in oldinfo:
+                        oldinfo.pop("all_prompts", None)
+                    if "all_negative_prompts" in oldinfo:
+                        oldinfo.pop("all_negative_prompts", None)
+                    if "infotexts" in oldinfo:
+                        oldinfo.pop("infotexts", None)
+                    response.info = json.dumps(oldinfo)
                     return response
+                elif req.task == 'upscale_from_feed':
+                    # only get the one image (in base64)
+                    intermediate_image = self.img2imgapi(req.img2img_payload).images
+                    print('finished intermediate img2img')
+                    try:
+                        # update the base64 image # note might need to change to req.extras_single_payload['image'] if this does not work
+                        req.extras_single_payload.image = intermediate_image[0]
+                        response = self.extras_single_image_api(req.extras_single_payload)
+                        response.image = self.post_invocations([response.image], quality)[0]
+                        response.parameters.clear()
+                        oldinfo = json.loads(response.info)
+                        if "all_prompts" in oldinfo:
+                            oldinfo.pop("all_prompts", None)
+                        if "all_negative_prompts" in oldinfo:
+                            oldinfo.pop("all_negative_prompts", None)
+                        if "infotexts" in oldinfo:
+                            oldinfo.pop("infotexts", None)
+                        response.info = json.dumps(oldinfo)
+                        # print(f"log@{datetime.datetime.now().strftime(f'%Y%m%d%H%M%S')} ### get_cmd_flags is {self.get_cmd_flags()}")
+                        return response
+                    except Exception as e:  # this is in fact obselete, because there will be a earlier return if OOM, won't reach here, but leaving here just in case
+                        print(
+                            f"An error occurred: {e}, step one upscale failed, reverting to just 4x upscale without Img2Img process")
                 elif req.task == 'extras-single-image':
                     response = self.extras_single_image_api(req.extras_single_payload)
                     response.image = self.post_invocations([response.image], quality)[0]
+                    if "info" in response:
+                        oldinfo = json.loads(response.info)
+                        if "all_prompts" in oldinfo:
+                            oldinfo.pop("all_prompts", None)
+                        if "all_negative_prompts" in oldinfo:
+                            oldinfo.pop("all_negative_prompts", None)
+                        if "infotexts" in oldinfo:
+                            oldinfo.pop("infotexts", None)
+                        response.info = json.dumps(oldinfo)
                     return response
                 elif req.task == 'extras-batch-images':
                     response = self.extras_batch_images_api(req.extras_batch_payload)
@@ -928,6 +1099,30 @@ def invocations(self, req: models.InvocationsRequest):
                 elif req.task == 'interrogate':
                     response = self.interrogateapi(req.interrogate_payload)
                     return response
+
+                elif req.task == 'get-progress':
+                    response = self.progressapi(req.progress_payload)
+                    print(response)
+                    return response
+                elif req.task == 'get-options':
+                    response = self.get_config()
+                    return response
+                elif req.task == 'get-SDmodels':
+                    response = self.get_sd_models()
+                    return response
+                elif req.task == 'get-upscalers':
+                    response = self.get_upscalers()
+                    return response
+                elif req.task == 'get-memory':
+                    response = self.get_memory()
+                    return response
+                elif req.task == 'get-cmd-flags':
+                    response = self.get_cmd_flags()
+                    return response
+                elif req.task == 'do-nothing':
+                    print("nothing has happened")
+                    return "nothing has happened"
+
                 elif req.task.startswith('/'):
                     if req.extra_payload:
                         response = requests.post(url=f'http://0.0.0.0:8080{req.task}', json=req.extra_payload)

modules/api/models.py

@@ -326,6 +326,9 @@ class InvocationsRequest(BaseModel):
     extras_batch_payload: Optional[ExtrasBatchImagesRequest]
     interrogate_payload: Optional[InterrogateRequest]
     extra_payload: Optional[dict]
+    user_input: Optional[dict]
+    progress_payload: Optional[ProgressRequest]
+    cn_x3_image: Optional[str]  # this is now obselete, but kept here in case of rollback
 
 
 class InvocationsErrorResponse(BaseModel):

modules/shared.py

@@ -159,8 +159,9 @@ def get_bucket_and_key(s3uri):
 def s3_download(s3uri, path):
     global cache
 
-    print('---path---', path)
-    os.system(f'ls -l {os.path.dirname(path)}')
+    #TODO: to delete
+    # print('---path---', path)
+    # os.system(f'ls -l {os.path.dirname(path)}')
 
     pos = s3uri.find('/', 5)
     bucket = s3uri[5 : pos]