建筑后端代码

import requests
import base64
import os
from PIL import Image
import io
import random


def encode_pil_to_base64(image):
    """
    Encode a PIL image to a Base64 string.
    """
    with io.BytesIO() as output_bytes:
        image.save(output_bytes, format="PNG")  # Using PNG format
        bytes_data = output_bytes.getvalue()
    return base64.b64encode(bytes_data).decode("utf-8")


def save_decoded_image(b64_image, folder_path, image_name):
    """
    Save a Base64 encoded image to a file in the specified folder.
    Automatically adjusts the filename by adding a sequence number
    if a file with the same name already exists.
    """
    # Initialize the sequence number and construct the initial file path
    seq = 0
    output_path = os.path.join(folder_path, f"{image_name}.png")

    # Check if the file exists and adjust the filename if necessary
    while os.path.exists(output_path):
        seq += 1
        # Construct a new filename with the sequence number
        output_path = os.path.join(folder_path, f"{image_name}({seq}).png")

    # Save the image to the new path
    with open(output_path, 'wb') as image_file:
        image_file.write(base64.b64decode(b64_image))
    print(f"Image saved to: {output_path}")


def main():
    while True:  # Infinite loop
        # Request the user to input the image file path
        user_input_path = input("Please enter the path of the reference image (or type 'exit' to quit): ")

        # Exit the loop if the user types 'exit'
        if user_input_path.lower() == 'exit':
            break

        if not os.path.exists(user_input_path) or not os.path.isfile(user_input_path):
            print("The specified path does not exist or is not a file. Please try again.")
            continue

        # Request the user to input the folder path for saving the generated image
        user_save_folder = input("Please enter the folder path where you'd like to save the generated image: ")

        if not os.path.exists(user_save_folder) or not os.path.isdir(user_save_folder):
            print("The specified save path does not exist or is not a folder. Please try again.")
            continue

        # Open and encode the reference image to Base64
        with Image.open(user_input_path) as img:
            encoded_image = encode_pil_to_base64(img)
            img_width, img_height = img.size  # Image dimensions

        # Generate a unique name for the saved image
        image_name = f"generated_image_{random.randint(1000, 9999)}"

        # API URL
        url = "http://127.0.0.1:7860/sdapi/v1/txt2img"  # Update to img2img API if needed

        # Construct the request payload
        data = {
            "prompt": "<lora:CWG_archisketch_v1:1>,Building,pre sketch,masterpiece,best quality,featuring markers,(3D:0.7)",
            "negative_prompt": "blurry, lower quality, glossy finish,insufficient contrast",
            "init_images": [encoded_image],  # Encoded image in a list
            "steps": 30,
            "width": img_width,
            "height": img_height,
            "seed": random.randint(1, 10000000),
            "alwayson_scripts": {
                "ControlNet": {
                    "args": [
                        {
                            "enabled": "true",
                            "pixel_perfect": "true",
                            "module": "canny",
                            "model": "control_v11p_sd15_canny_fp16 [b18e0966]",
                            "weight": 1,
                            "image": encoded_image
                        },
                        {
                            "enabled": "true",
                            "pixel_perfect": "true",
                            "module": "depth",
                            "model": "control_v11f1p_sd15_depth_fp16 [4b72d323]",
                            "weight": 1,
                            "image": encoded_image
                        }
                    ]
                }
            }
        }

        # Send the request and get the response
        response = requests.post(url, json=data)
        response_json = response.json()

        # Save the generated image to the specified folder with a unique name
        save_decoded_image(response_json['images'][0], user_save_folder, image_name)


if __name__ == '__main__':
    main()