xcfg

main.py

@@ -34,10 +34,13 @@ def decode_and_save_base64(base64_str, save_path):
     with open(save_path, "wb") as file:
         file.write(base64.b64decode(base64_str))
 
+
+
 def parse_input(input_string):
     # Set default payload values
     default_payload = {
         "prompt": "",
+        "seed": -1, # Random seed
         "negative_prompt": "ugly, bad face, distorted",
         "enable_hr": False,
         "Sampler": "DPM++ SDE Karras",
@@ -58,10 +61,16 @@ def parse_input(input_string):
     matches = re.finditer(r"(\w+):", input_string)
     last_index = 0
 
-    # Initialize script_args with default values and placeholder enums
     script_args = [0, "", [], 0, "", [], 0, "", [], True, False, False, False, False, False, False, 0, False]
     script_name = None
 
+
+    script_args = [0, "", [], 0, "", [], 0, "", [], True, False, False, False, False, False, False, 0, False]
+    script_name = None
+
+    slot_mapping = {0: (0, 1), 1: (3, 4), 2: (6, 7)}
+    slot_index = 0
+
     for match in matches:
         key = match.group(1).lower()
         value_start_index = match.end()
@@ -78,27 +87,34 @@ def parse_input(input_string):
             key = "denoising_strength"
         if key == "ng":
             key = "negative_prompt"
+        if key == "cfg":
+            key = "cfg_scale"
 
         if key in default_payload:
             payload[key] = value
-        elif key in ["xsr", "xsteps", "xds"]:
+        elif key in ["xsr", "xsteps", "xds", "xcfg", "nl", "ks", "rs"]:
             script_name = "x/y/z plot"
-            if key == "xsr":
-                script_args[0] = 7  # Enum value for xsr
-                script_args[1] = value
-            elif key == "xsteps":
-                try:
-                    steps_values = [int(x) for x in value.split(',')]
-                    if all(1 <= x <= 70 for x in steps_values):
-                        script_args[3] = 4  # Enum value for xsteps
-                        script_args[4] = value
-                    else:
-                        raise ValueError("xsteps values must be between 1 and 70.")
-                except ValueError:
-                    raise ValueError("xsteps must contain only integers.")
-            elif key == "xds":
-                script_args[6] = 22  # Enum value for xds
-                script_args[7] = value
+            if slot_index < 3:
+                script_slot = slot_mapping[slot_index]
+                if key == "xsr":
+                    script_args[script_slot[0]] = 7  # Enum value for xsr
+                    script_args[script_slot[1]] = value
+                elif key == "xsteps":
+                    script_args[script_slot[0]] = 4  # Enum value for xsteps
+                    script_args[script_slot[1]] = value
+                elif key == "xds":
+                    script_args[script_slot[0]] = 22  # Enum value for xds
+                    script_args[script_slot[1]] = value
+                elif key == "xcfg":
+                    script_args[script_slot[0]] = 6  # Enum value for CFG Scale
+                    script_args[script_slot[1]] = value
+                slot_index += 1
+            elif key == "nl":
+                script_args[9] = False  # Draw legend
+            elif key == "ks":
+                script_args[10] = True  # Keep sub images
+            elif key == "rs":
+                script_args[11] = True  # Set random seed to sub images
         else:
             prompt.append(f"{key}:{value}")
 
@@ -113,11 +129,27 @@ def parse_input(input_string):
         payload["script_args"] = script_args
 
     return payload
-
-
 def create_caption(payload, user_name, user_id, info):
     caption = f"**[{user_name}](tg://user?id={user_id})**\n\n"
     prompt = payload["prompt"]
+    print(payload["prompt"])
+    print(info)
+    # Steps: 3, Sampler: Euler, CFG scale: 7.0, Seed: 4094161400, Size: 512x512, Model hash: 15012c538f, Model: realisticVisionV60B1_v51VAE, Denoising strength: 0.35, Version: v1.8.0-1-g20cdc7c
+
+    # Define a regular expression pattern to match the seed value
+    seed_pattern = r"Seed: (\d+)"
+
+    # Search for the pattern in the info string
+    match = re.search(seed_pattern, info)
+
+    # Check if a match was found and extract the seed value
+    if match:
+        seed_value = match.group(1)
+        print(f"Seed value: {seed_value}")
+        caption += f"**{seed_value}**\n"
+    else:
+        print("Seed value not found in the info string.")
+
     caption += f"**{prompt}**\n"
 
     if len(caption) > 1024:
@@ -188,7 +220,11 @@ def img2img(client, message):
         return
 
     payload = parse_input(msgs[1])
+    print(f"input:\n{payload}")
     photo = message.reply_to_message.photo
+    # prompt_from_reply = message.reply_to_message.
+    # orginal_prompt = app.reply_to_message.message
+    # print(orginal_prompt)
     photo_file = app.download_media(photo)
     init_image = encode_file_to_base64(photo_file)
     os.remove(photo_file)  # Clean up downloaded image file