Spaces:
Sleeping
Sleeping
File size: 3,896 Bytes
4f48282 |
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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
# img_gen.py
import sys
import os
import random
from huggingface_hub import InferenceClient
from datetime import datetime
from config.config import models, prompts, api_token # Direct import
import modal
# Define the Modal image
image = (
modal.Image.from_registry("nvidia/cuda:12.2.0-devel-ubuntu22.04", add_python="3.9")
.apt_install(
"git",
)
.pip_install(
"diffusers",
"transformers",
"torch",
"accelerate",
"gradio>=4.44.1",
"safetensors",
"pillow",
"sentencepiece",
"hf_transfer",
"huggingface_hub[hf_transfer]",
"aria2", # aria2 for ultra-fast parallel downloads
)
.env(
{
"HF_HUB_ENABLE_HF_TRANSFER": "1", "HF_HOME": "HF_HOME"
}
)
)
# Create a Modal app
app = modal.App("img-gen-modal", image=image)
with image.imports():
import diffusers
import os
import gradio
import torch
import sentencepiece
import transformers
from huggingface_hub import InferenceClient, login
@app.function(
secrets=[modal.Secret.from_name("huggingface-token")],
gpu="t4",
timeout=600
)
def generate_image(prompt_alias, team_color, model_alias, custom_prompt, height=360, width=640, num_inference_steps=20, guidance_scale=2.0, seed=-1):
#from huggingface_hub import InferenceClient
# 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."
# 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)
# Print the formatted prompt for debugging
print("\nFormatted Prompt:")
print(prompt)
# Append the custom prompt (if provided)
if custom_prompt and len(custom_prompt.strip()) > 0:
prompt += " " + custom_prompt.strip()
# Randomize the seed if needed
if seed == -1:
seed = random.randint(0, 1000000)
# HF LOGIN
print("Initializing HF TOKEN")
hf_token = os.environ["HF_TOKEN"]
print(hf_token)
print("HF TOKEN:")
login(token=hf_token)
print("model_name:")
print(model_name)
# Initialize the InferenceClient
try:
client = InferenceClient(model_name, token=api_token)
except Exception as e:
return None, f"ERROR: Failed to initialize InferenceClient. Details: {e}"
#Generate the image
try:
image = client.text_to_image(
prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
width=width,
height=height,
seed=seed
)
except Exception as e:
return None, f"ERROR: Failed to generate image. Details: {e}"
# Save the image with a timestamped filename
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_filename = f"{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}"
return output_filename, "Image generated successfully!" |