建筑转手稿

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, base_path):
    """
    Save a Base64 encoded image to a file.
    Automatically adjusts the filename by adding a sequence number
    if a file with the same name already exists.
    """
    # Initialize the sequence number and the file path
    seq = 0
    output_path = base_path

    # Check if the file exists and adjust the filename if necessary
    while os.path.exists(output_path):
        seq += 1
        # Split the base path into the directory, basename, and extension
        dir_name, base_name = os.path.split(base_path)
        name, ext = os.path.splitext(base_name)
        # Construct a new filename with the sequence number
        output_path = os.path.join(dir_name, f"{name}({seq}){ext}")

    # 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():
    # API URL
    url = "http://127.0.0.1:7860/sdapi/v1/txt2img"  # Update to img2img API if needed

    # Path to the folder containing reference images
    reference_image_dir = r"C:\Users\wujie1\Desktop\图片参考素材"

    # Find the latest file in the reference image folder
    latest_file_path = max(
        [os.path.join(reference_image_dir, f) for f in os.listdir(reference_image_dir) if os.path.isfile(os.path.join(reference_image_dir, f))],
        key=os.path.getmtime,
        default=None  # Default to None if the folder is empty
    )

    if latest_file_path is None:
        print("No image files found in the specified folder.")
        return

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

    # Construct the request payload
    data = {
        "prompt": "<lora:CWG_archisketch_v1:1>,Building,masterpiece,best quality,pre sktch,",
        "negative_prompt": "blurry, lower quality, 3D",
        "init_images": [encoded_image],  # Encoded image in a list
        "steps": 25,
        "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()

    # Define the base path for saving the image (without sequence number)
    base_save_path = r"C:\Users\wujie1\Downloads\Generated_Image.png"
    save_decoded_image(response_json['images'][0], base_save_path)

if __name__ == '__main__':
    main()