image-to-image.py

import base64
import os
import requests
from dotenv import load_dotenv
from PIL import Image, ImageOps

# Load environment variables
load_dotenv()

def calculate_best_fit_dimension(original_width, original_height, target_dimensions):
    original_aspect_ratio = original_width / original_height
    best_fit = None
    smallest_diff = float('inf')
    for width, height in target_dimensions:
        target_aspect_ratio = width / height
        aspect_diff = abs(target_aspect_ratio - original_aspect_ratio)
        if aspect_diff < smallest_diff:
            best_fit = (width, height)
            smallest_diff = aspect_diff
    return best_fit

def resize_to_best_fit(input_path, output_path, target_dimensions):
    with Image.open(input_path) as img:
        best_fit_dimension = calculate_best_fit_dimension(img.width, img.height, target_dimensions)
        # Resize the image to fill the best fit dimension
        img = img.resize(best_fit_dimension, Image.Resampling.LANCZOS)
        img.save(output_path)


engine_id = "stable-diffusion-xl-1024-v1-0"
api_host = os.getenv("API_HOST", "https://api.stability.ai")
api_key = os.getenv("STABILITY_API_KEY")

if api_key is None:
    raise Exception("Missing Stability API key.")

# Define the target dimensions as per API requirement
target_dimensions = [
    (1024, 1024), (1152, 896), (1216, 832),
    (1344, 768), (1536, 640), (640, 1536),
    (768, 1344), (832, 1216), (896, 1152)
]

# Resize the image
input_image_path = "franck.jpg"
output_image_path = "franck_resized.jpg"
resize_to_best_fit(input_image_path, output_image_path, target_dimensions)

# Make the API request using the resized image
response = requests.post(
    f"{api_host}/v1/generation/{engine_id}/image-to-image",
    headers={
        "Accept": "application/json",
        "Authorization": f"Bearer {api_key}"
    },
    files={
        "init_image": open(output_image_path, "rb")
    },
    data={
        "image_strength": 0.35,
        "init_image_mode": "IMAGE_STRENGTH",
        "text_prompts[0][text]": "Add some colors to the image, make it modern.",
        "cfg_scale": 7,
        "samples": 1,
        "steps": 30,
    }
)

if response.status_code != 200:
    raise Exception("Non-200 response: " + str(response.text))

data = response.json()

# Save the transformed images
for i, image in enumerate(data["artifacts"]):
    with open(f"franck_transformed_{i}.jpg", "wb") as f:
        f.write(base64.b64decode(image["base64"]))