Spaces:
Sleeping
Sleeping
| """ 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" | |