equilink-site/server/prompt_chaining/step_4_create_profile_images.py

183 lines
7.2 KiB
Python
Raw 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
import asset_generation.images_leonardo as images_leonardo
from models.gemini_model import GeminiModel
dotenv.load_dotenv()
def create_images(ai_model, prompt, master_file_path, num_images):
# 4.1 Setup model details
model_details = json.loads(open("asset_generation/leonard_anime_styles.json", "r").read())
model_id = model_details["models"][0]["modelId"]
anime_general_style = next(
style for style in model_details["models"][0]["styles"]
if style["style name"] == "Anime General"
)
# 4.2 Setup style UUID
style_uuid = anime_general_style["styleUUID"]
# 4.3 Setup number of images
num_images = 4
# 4.4 Create image descriptions
# If prompt is None, we need to create image descriptions automatically
if prompt is None or prompt == "":
# # 4.4 Create image descriptions
# image_descriptions = image_post_description(ai_model, master_file_path, num_images)
# # 4.5 loop through image descriptions and create profile images
# for image_description in image_descriptions["profile_image_descriptions"]:
# print(image_description["description"] + "\n")
# # 4.6 Create profile images
# create_profile_images(image_description, master_file_path, model_id, style_uuid, num_images, consistent=True, max_retries=10, delay=5)
# Simple fix for now is to feed in the agents concept into the prompt
agent_master_json = None
with open(master_file_path, 'r', encoding='utf-8') as file:
agent_master_json = json.load(file)
prompt = agent_master_json['agent']['concept']
create_profile_images(prompt, master_file_path, model_id, style_uuid, num_images, consistent=False, max_retries=10, delay=5)
# 4.7 Create profile images based on users prompt
else:
create_profile_images(prompt, master_file_path, model_id, style_uuid, num_images, consistent=False, max_retries=10, delay=5)
def image_post_description(ai_model, master_file_path, num_images):
manager = ContentGenerator()
image_descriptions = None
# step 4.2: load the agent json file
agent_master_json = None
with open(master_file_path, 'r', encoding='utf-8') as file:
agent_master_json = json.load(file)
# step 4.3: extract agent from master json
agent_details = agent_master_json['agent']['agent_details']
print("Crafting prompt for AI to create a new profile image descriptions")
prompt_4_vars = {
"agent_name": agent_details["name"],
"agent_json": json.dumps(agent_details),
"number_of_images": num_images
}
# step 2.4: Run the prompt
print("Sending prompt to AI to create a new profile image descriptions")
image_descriptions = manager.run_prompt(
prompt_key="prompt_4 (Agent Profile Image)",
template_vars=prompt_4_vars,
ai_model=ai_model,
)
# print(f"image_descriptions is: {image_descriptions}")
return image_descriptions
#--------------------------------
# Step 4: Generate a number of images for profile image
#--------------------------------
def create_profile_images(prompt, master_file_path, model_id, style_uuid, num_images, consistent=False, max_retries=10, delay=5):
# step 4.1: load existing profile image options from templates directory
manager = ContentGenerator()
profile_image_options_template = manager.create_new_template_json(TemplateType.PROFILE_IMAGE_OPTIONS)
profile_image_template = manager.create_new_template_json(TemplateType.PROFILE_IMAGE)
# step 4.2: load the agent json file
agent_master_json = None
with open(master_file_path, 'r', encoding='utf-8') as file:
agent_master_json = json.load(file)
# step 4.3: extract agent from master json
agent_details = agent_master_json['agent']['agent_details']
# step 4.3.1: check to make sure we have a profile image placeholder
if agent_master_json['agent']['profile_image'] is None or agent_master_json['agent']['profile_image'] == {}:
agent_master_json['agent']['profile_image'] = profile_image_template['profile_image']
# step 4.4: generate the image
if consistent:
response = images_leonardo.generate_inconsistent_image_lambda(prompt, model_id, style_uuid, num_images)
else:
response = images_leonardo.generate_inconsistent_image_lambda(prompt, model_id, style_uuid, num_images)
if not response['generations_by_pk']['generated_images']:
raise Exception("No Images Generated")
generations_by_pk = response['generations_by_pk'].get("generated_images")
generated_url = generations_by_pk[0].get('url')
generated_id = generations_by_pk[0].get('id')
# step 4.6: append new URL to the array
profile_image_options_template["profile_image_options"].append(response)
profile_image_template["profile_image"]["details"]["generationId"] = generated_id
profile_image_template["profile_image"]["details"]["url"] = generated_url
profile_image_template["profile_image"]["details"]["image_id"] = generated_id
# step 4.7: create the save path, use 0 as we are not in a season or episode
save_path = manager.create_filepath(agent_details["name"], 0, 0, TemplateType.PROFILE_IMAGE_OPTIONS)
save_path_profile_image = manager.create_filepath(agent_details["name"], 0, 0, TemplateType.PROFILE_IMAGE)
# step 4.8: save the profile image options to a file
manager.save_json_file(save_path, profile_image_options_template)
# if this is consistent, we will multiple genreation ids to track. So its better to store the generation id of the image the user picks.
if not consistent:
manager.save_json_file(save_path_profile_image, profile_image_template)
# step 4.9: append to master file
agent_master_json = manager.append_profile_image_options(agent_master_json, profile_image_options_template)
agent_master_json = manager.append_profile_image(agent_master_json, profile_image_template)
# step 4.10: save the master data to a file
print("Saving the master data to a file")
manager.save_json_file(
save_path=master_file_path,
json_data=agent_master_json
)
if __name__ == "__main__":
# 4.1 Setup number of images
num_images = 4
# 4.2 Setup AI model
ai_model = GeminiModel()
master_file_path = "configs/CipherCat/CipherCat_master.json"
# 4.2.1 Setup prompt
# prompt = "a pygmy marmoset who owns 30% of trumps meme coin"
prompt = ""
# 4.3 Create profile images
create_images(ai_model, prompt, master_file_path, num_images)