Spaces:
Running
Running
# img_gen.py | |
import sys | |
import os | |
import random | |
from huggingface_hub import InferenceClient, login | |
from datetime import datetime | |
from config.config import prompts, api_token | |
from config.models import models | |
from metadata.metadata import fetch_metadata | |
def generate_image( | |
adventurer_id, | |
#prompt_alias, | |
custom_prompt, | |
characer_dropdown, | |
model_alias, | |
height=360, | |
width=640, | |
num_inference_steps=20, | |
guidance_scale=2.0, | |
seed=-1): | |
adventurer = fetch_metadata(adventurer_id) | |
print(f"ANDRE {adventurer['name']}") | |
id = adventurer['id'] | |
name = adventurer['name'] | |
weapon_equipment = adventurer['weapon'] | |
head_equipment = adventurer['head'] | |
hand_equipment = adventurer['hand'] | |
chest_equipment = adventurer['chest'] | |
waist_equipment = adventurer['waist'] | |
foot_equipment = adventurer['foot'] | |
gold_equipment = adventurer['gold'] | |
beast_last_battle = adventurer['beast'] | |
# Set the custom prompt variables | |
if characer_dropdown == "Adventurer Portait": | |
prompt = f"A portait of a medieval, fantasy adventurer, holding only a {weapon_equipment} (depending on his weapon, make the characer dressed as a warrior, or as a hunter or as a wizard). He is also equiped in the head with a {head_equipment}, the hands with {hand_equipment}, the chest with a {chest_equipment}, and the waist with a {waist_equipment}. Please be sure to use only medieval items that were possble to be made in that period. Make the adventurer unshaved and dirty, he have been figthing for days down into the dungeons. Unreal Engine render style, photorealistic, atmospheric light, realistic fantasy style." | |
if characer_dropdown == "Beast Portait": | |
prompt = f"A massive {beast_last_battle} stands in its natural domain—a place that reflects its primal essence. The beast's eyes burn with an eerie, intelligent light, and its powerful limbs are poised for action, as if ready to pounce or defend its territory. The air is thick with tension, and the ground is littered with signs of its dominance—broken terrain, scattered bones, or remnants of its prey. In the background, a dramatic landscape unfolds: dense forests, jagged mountains, dark caves, or misty swamps. The atmosphere is both terrifying and awe-inspiring, capturing the essence of a fantasy world where magic and monsters reign supreme. Unreal Engine render style, photorealistic, atmospheric light, realistic fantasy style." | |
if characer_dropdown == "Last Battle": | |
prompt = f"A battle between a medieval fantasy adventurer, and a massive {beast_last_battle}. The adventurer is holding only a {weapon_equipment}. He is also equiped in the head with {head_equipment}, the hands with {hand_equipment}, the chest with {chest_equipment}, the waist with {waist_equipment}, the feet with {foot_equipment}. Use only medieval items that were possble to be made in that period. The scene location is the natural ambient of the {beast_last_battle}. Unreal Engine render style, photorealistic, realistic fantasy style." | |
elif characer_dropdown == "Loot Bag": | |
prompt = f"A loot bag from a medieval fantasy adventurer and his equipments. On the floor also a {weapon_equipment} a {head_equipment}, a {hand_equipment}, a {chest_equipment}, a {waist_equipment}, and a {foot_equipment}. Please sure to use only medieval items that were possble to be made in that period. Inside the bag also {gold_equipment} gold coins. Atmospheric light, cavern, dungeon context. Unreal Engine render style, photorealistic, realistic fantasy style." | |
else: | |
pass | |
# Find the selected prompt and model | |
try: | |
#prompt = next(p for p in prompts if p["alias"] == prompt_alias)["text"] | |
model_name = next(m for m in models if m["alias"] == model_alias)["name"] | |
except StopIteration: | |
return None, "ERROR: Invalid prompt or model selected." | |
# Print the original prompt and dynamic values for debugging | |
print("Original Prompt:") | |
print(prompt) | |
# Append the custom prompt (if provided) | |
if custom_prompt and len(custom_prompt.strip()) > 0: | |
prompt += " " + custom_prompt.strip() | |
# Print the formatted prompt for debugging | |
print("\nFormatted Prompt:") | |
print(prompt) | |
# Randomize the seed if needed | |
if seed == -1: | |
seed = random.randint(0, 1000000) | |
# HF LOGIN | |
print("Initializing HF TOKEN") | |
print (api_token) | |
# login(token=api_token) | |
# print("model_name:") | |
# print(model_name) | |
# Initialize the InferenceClient | |
try: | |
print("-----INITIALIZING INFERENCE-----") | |
client = InferenceClient(model_name, token=api_token) | |
print("Inference activated") | |
except Exception as e: | |
return None, f"ERROR: Failed to initialize InferenceClient. Details: {e}" | |
#Generate the image | |
try: | |
print("-----GENERATING IMAGE-----") | |
print("-----HOLD ON-----") | |
image = client.text_to_image( | |
prompt, | |
guidance_scale=guidance_scale, | |
num_inference_steps=num_inference_steps, | |
width=width, | |
height=height, | |
seed=seed | |
) | |
print("-----IMAGE GENERATED SUCCESSFULLY!-----") | |
except Exception as e: | |
return None, f"ERROR: Failed to generate image. Details: {e}" | |
# Save the image with a timestamped filename | |
print("-----SAVING-----", image) | |
path = "images" | |
message = f"Image generated successfully! Call the banners! \n\n=====ADVENTURER GENERATED===== \nID: {id}\nNAME: {name}\nWEAPON: {weapon_equipment}\nHEAD: {head_equipment}\nHAND: {hand_equipment}\nCHEST: {chest_equipment}\nWAIST: {waist_equipment}\nLAST BATTLE BEAST: {beast_last_battle}" | |
file_name_extension = f"{id} {name} {weapon_equipment} {head_equipment} {hand_equipment} {chest_equipment} {waist_equipment} {beast_last_battle}" | |
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
output_filename = f"{path}/{timestamp}_{seed}_{model_alias.replace(' ', '_').lower()}_{characer_dropdown.replace(' ', '_').lower()}_{file_name_extension.replace(' ', '_').lower()}.png" | |
# output_filename = f"{path}/{timestamp}_{seed}_{model_alias.replace(' ', '_').lower()}_{prompt_alias.replace(' ', '_').lower()}_{characer_dropdown.replace(' ', '_').lower()}_{file_name_extension.replace(' ', '_').lower()}.png" | |
try: | |
image.save(output_filename) | |
except Exception as e: | |
return None, f"ERROR: Failed to save image. Details: {e}" | |
print("-----DONE!-----") | |
print("-----CALL THE BANNERS!-----") | |
return output_filename, message |