201 lines
5.5 KiB
Python
201 lines
5.5 KiB
Python
![]() |
#
|
||
|
# Module: template_types
|
||
|
#
|
||
|
# This module implements the TemplateType class for defining the type of template.
|
||
|
#
|
||
|
# Title: Template Types
|
||
|
# Summary: Template types implementation.
|
||
|
# Authors:
|
||
|
# - @TheBlockRhino
|
||
|
# Created: 2024-12-31
|
||
|
# Last edited by: @TheBlockRhino
|
||
|
# Last edited date: 2025-01-12
|
||
|
# URLs:
|
||
|
# - https://arai-ai.io
|
||
|
# - https://github.com/ARAI-DevHub/arai-ai-agents
|
||
|
# - https://x.com/TheBlockRhino
|
||
|
|
||
|
import requests
|
||
|
import os
|
||
|
import dotenv
|
||
|
import time
|
||
|
import json
|
||
|
import sys
|
||
|
|
||
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
|
||
|
# custom ARAI imports
|
||
|
from utils.content_generator import ContentGenerator
|
||
|
from utils.template_types import TemplateType
|
||
|
|
||
|
dotenv.load_dotenv()
|
||
|
|
||
|
# Define the Lambda endpoint URL
|
||
|
get_inconsistent_image_lambda_url = "https://46i9cnowhh.execute-api.us-east-1.amazonaws.com/getImageInconsistent"
|
||
|
|
||
|
# Function to call the AWS Lambda
|
||
|
def generate_inconsistent_image_lambda(prompt, model_id, style_uuid, num_images):
|
||
|
|
||
|
url = get_inconsistent_image_lambda_url
|
||
|
|
||
|
headers = {
|
||
|
"Content-Type": "application/json"
|
||
|
}
|
||
|
|
||
|
payload = {
|
||
|
"prompt": prompt,
|
||
|
"modelId": model_id,
|
||
|
"styleUUID": style_uuid,
|
||
|
"num_images": num_images
|
||
|
}
|
||
|
|
||
|
try:
|
||
|
# Send the POST request with the payload
|
||
|
# Use dumpts to convert the payload to a string
|
||
|
response = requests.post(url, headers=headers, data=json.dumps(payload))
|
||
|
|
||
|
# Check for HTTP errors
|
||
|
if not response.ok:
|
||
|
print(f"HTTP error! Status: {response.status_code}, Response: {response.text}")
|
||
|
response.raise_for_status()
|
||
|
|
||
|
# Parse and return the response JSON
|
||
|
data = response.json()
|
||
|
return data
|
||
|
|
||
|
except requests.RequestException as error:
|
||
|
print("Error calling Lambda:", error)
|
||
|
raise error
|
||
|
|
||
|
#--------------------------------
|
||
|
# Generate multiple images from a prompt that are inconsistent with each other
|
||
|
#--------------------------------
|
||
|
def generated_image_inconsistent(prompt, model_id, style_uuid, num_images):
|
||
|
url = "https://cloud.leonardo.ai/api/rest/v1/generations"
|
||
|
|
||
|
payload = {
|
||
|
"modelId": model_id,
|
||
|
"presetStyle": "DYNAMIC",
|
||
|
"scheduler": "LEONARDO",
|
||
|
"sd_version": "SDXL_LIGHTNING",
|
||
|
"contrast": 1.3,
|
||
|
"prompt": prompt,
|
||
|
"num_images": num_images,
|
||
|
"width": 512,
|
||
|
"height": 512,
|
||
|
"alchemy": True,
|
||
|
"styleUUID": style_uuid,
|
||
|
"enhancePrompt": False,
|
||
|
"nsfw": True,
|
||
|
"public": False,
|
||
|
}
|
||
|
headers = {
|
||
|
"accept": "application/json",
|
||
|
"content-type": "application/json",
|
||
|
"authorization": "Bearer " + os.getenv("LEONARDO_API_KEY", "")
|
||
|
}
|
||
|
|
||
|
try:
|
||
|
response = requests.post(url, json=payload, headers=headers)
|
||
|
return response.json(), payload
|
||
|
except Exception as e:
|
||
|
print("Error:", str(e))
|
||
|
return None, payload
|
||
|
|
||
|
#--------------------------------
|
||
|
# Generate multiple images from a prompt that are consistent with each other
|
||
|
#--------------------------------
|
||
|
def generated_image_consistent(prompt, model_id, style_uuid, num_images):
|
||
|
url = "https://cloud.leonardo.ai/api/rest/v1/generations"
|
||
|
|
||
|
payload = {
|
||
|
"modelId": model_id,
|
||
|
"presetStyle": "DYNAMIC",
|
||
|
"scheduler": "LEONARDO",
|
||
|
"sd_version": "SDXL_LIGHTNING",
|
||
|
"contrast": 1.3,
|
||
|
"prompt": prompt,
|
||
|
"num_images": num_images,
|
||
|
"width": 1024,
|
||
|
"height": 1024,
|
||
|
"alchemy": True,
|
||
|
"styleUUID": style_uuid,
|
||
|
"enhancePrompt": False,
|
||
|
"nsfw": True,
|
||
|
"public": False,
|
||
|
"num_inference_steps": 10,
|
||
|
"guidance_scale": 7,
|
||
|
|
||
|
}
|
||
|
headers = {
|
||
|
"accept": "application/json",
|
||
|
"content-type": "application/json",
|
||
|
"authorization": "Bearer " + os.getenv("LEONARDO_API_KEY", "")
|
||
|
}
|
||
|
|
||
|
try:
|
||
|
response = requests.post(url, json=payload, headers=headers)
|
||
|
return response.json(), payload
|
||
|
except Exception as e:
|
||
|
print("Error:", str(e))
|
||
|
return None, payload
|
||
|
|
||
|
#--------------------------------
|
||
|
# Get the image URL from a generation ID
|
||
|
#--------------------------------
|
||
|
def get_image_url(generation_id):
|
||
|
url = "https://cloud.leonardo.ai/api/rest/v1/generations/" + generation_id
|
||
|
|
||
|
headers = {
|
||
|
"accept": "application/json",
|
||
|
"authorization": "Bearer " + os.getenv("LEONARDO_API_KEY", "")
|
||
|
}
|
||
|
|
||
|
response = requests.get(url, headers=headers)
|
||
|
|
||
|
return response.json()
|
||
|
|
||
|
#--------------------------------
|
||
|
# Get list of models
|
||
|
#--------------------------------
|
||
|
def get_models():
|
||
|
url = "https://cloud.leonardo.ai/api/rest/v1/platformModels"
|
||
|
|
||
|
headers = {
|
||
|
"accept": "application/json",
|
||
|
"authorization": "Bearer " + os.getenv("LEONARDO_API_KEY", "")
|
||
|
}
|
||
|
|
||
|
response = requests.get(url, headers=headers)
|
||
|
|
||
|
return response.json()
|
||
|
|
||
|
# Get a list of model styles
|
||
|
#--------------------------------
|
||
|
def get_model_styles(model_id):
|
||
|
url = "https://cloud.leonardo.ai/api/rest/v1/platformModels/" + model_id + "/styles"
|
||
|
|
||
|
headers = {
|
||
|
"accept": "application/json",
|
||
|
"authorization": "Bearer " + os.getenv("LEONARDO_API_KEY", "")
|
||
|
}
|
||
|
|
||
|
response = requests.get(url, headers=headers)
|
||
|
|
||
|
return response.json()
|
||
|
|
||
|
#--------------------------------
|
||
|
# Get list of elements
|
||
|
#--------------------------------
|
||
|
def get_elements():
|
||
|
url = "https://cloud.leonardo.ai/api/rest/v1/elements"
|
||
|
|
||
|
headers = {
|
||
|
"accept": "application/json",
|
||
|
"authorization": "Bearer " + os.getenv("LEONARDO_API_KEY", "")
|
||
|
}
|
||
|
|
||
|
response = requests.get(url, headers=headers)
|
||
|
|
||
|
return response.json()
|