201 lines
5.5 KiB
Python
Raw Permalink Normal View History

2025-02-12 17:38:06 +05:30
#
# 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()