Spaces:
Sleeping
Sleeping
File size: 3,584 Bytes
af7831d 472af83 e1fd06b 4f48282 e1fd06b 959fd70 ef908f6 4f48282 af7831d 4f48282 af7831d e1fd06b 4f48282 e1fd06b f2f08ce e1fd06b c3e503f e19588a c3e503f e19588a c3e503f e19588a e1fd06b 4f48282 e1fd06b 4f48282 e1fd06b 4f48282 e1fd06b 4f48282 737c5f6 e1fd06b 737c5f6 4f48282 e1fd06b 4f48282 4eb4c43 e1fd06b 4f48282 e1fd06b 4f48282 e1fd06b 4f48282 e1fd06b caba2dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
# img_gen.py
import sys
import os
import random
from huggingface_hub import InferenceClient, login
from datetime import datetime
from config.config import models, prompts, api_token # Direct import
def generate_image(
prompt_alias,
team_color,
custom_prompt,
model_alias="FLUX.1-dev",
width=640,
height=360,
num_inference_steps=20,
guidance_scale=2.0,
seed=-1):
# LOAD CSS
try:
with open("src/aaa.css", "r") as f:
custom_css = f.read()
except FileNotFoundError:
print("Error: aaa.css not found!")
custom_css = "" # Or provide default CSS
# 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"]
model_name = f"black-forest-labs/{model_alias}"
except StopIteration:
return None, "ERROR: Invalid prompt or model selected."
# Determine the enemy color
enemy_color = "blue" if team_color.lower() == "red" else "red"
# if team.lower() == "red":
# winning_team_text = " The winning army is dressed in red armor and banners."
# elif team.lower() == "blue":
# winning_team_text = " The winning army is dressed in blue armor and banners."
# Print the original prompt and dynamic values for debugging
print("Original Prompt:")
print(prompt)
print(f"Enemy Color: {enemy_color}")
print(f"Team Color: {team_color.lower()}")
prompt = prompt.format(team_color=team_color.lower(), enemy_color=enemy_color)
# 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"
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_filename = f"{path}/{timestamp}_{model_alias.replace(' ', '_').lower()}_{prompt_alias.replace(' ', '_').lower()}_{team_color.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, "Image generated successfully! Call the banners!" |