diff --git a/main.py b/main.py index 82f989c..3e61f4e 100644 --- a/main.py +++ b/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