midjourney_generate_img.py

""" Generate an image using the Midjourney API"""
import io
import random
import requests
from PIL import Image
from progress_bar import print_progress_bar
from config import GOAPIKEY

IMAGINE_ENDPOINT= "https://api.midjourneyapi.xyz/mj/v2/imagine"
FETCH_ENDPOINT = "https://api.midjourneyapi.xyz/mj/v2/fetch"

headers = {
    "X-API-KEY": GOAPIKEY
}

def midjourney_generate_img(prompt, api_key=GOAPIKEY):
    """Generate an image using the Midjourney API

    Keyword arguments:
    prompt -- The prompt to generate the image from
    Return: An image saved in a .png file
    """
    headers["X-API-KEY"] = api_key
    img_generation_data = {
        "prompt": prompt,
        "aspect_ratio": "16:9",
        "process_mode": "fast",
        "webhook_endpoint": "",
        "webhook_secret": ""
    }
    create_img_response = requests.post(
        IMAGINE_ENDPOINT,
        headers=headers,
        json=img_generation_data,
        timeout=30)
    if create_img_response.status_code == 200:
        print("Request for an img to Midjourney: successfully!")
        task_id = create_img_response.json()['task_id']
    else:
        print(f"Image creation failed, please review details: {create_img_response.status_code} \
              {create_img_response.text}")

    print_progress_bar(50, msg='Generating MidJourney img, please wait...', bar_length=20)
    fetch_img_response = requests.post(FETCH_ENDPOINT,
                                       headers=headers,
                                       json={"task_id": task_id},
                                       timeout=30)
    status_img = fetch_img_response.json()['status']
    while status_img != "finished":
        # pause 10s
        print_progress_bar(10, msg='Generating MidJourney img, please wait...', bar_length=20)
        fetch_img_response = requests.post(FETCH_ENDPOINT,
                                           headers=headers,
                                           json={"task_id": task_id},
                                           timeout=30)
        status_img = fetch_img_response.json()['status']
        if status_img == "failed":
            print(
                f"Image generation failed, please review details: {fetch_img_response.status_code} \
                  {fetch_img_response.text}"
                  )

    # download task_result image
    print("Saving img...")
    task_result_image_url = fetch_img_response.json()['task_result']['image_url']
    image_response = requests.get(task_result_image_url, timeout=30)
    # saving img to output_img as png
    img = Image.open(io.BytesIO(image_response.content))
    img.save('output_img/midjourney_generated_img.png')
    print("Image saved in output_img/midjourney_generated_img.png")

    # divide img by 4 and save only one part
    img_width, img_height = img.size
    target_width = img_width // 2
    target_height = img_height // 2
    part = random.randint(1, 4) # select a random part
    if part == 1:
        img_cropped = img.crop((0, 0, target_width, target_height))  # Superior izquierda
    elif part == 2:
        img_cropped = img.crop((target_width, 0, img_width, target_height))  # Superior derecha
    elif part == 3:
        img_cropped = img.crop((0, target_height, target_width, img_height))  # Inferior izquierda
    else:
        img_cropped = img.crop((target_width, target_height, img_width, img_height))  # Inferior derecha
    # save the selected img
    img_cropped.save('output_img/midjourney_single_img.png')
    print("Single image saved in output_img/midjourney_single_img.png")
    return "output_img/midjourney_single_img.png"