Spaces:
Sleeping
Sleeping
File size: 4,633 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 132 133 134 |
#img_gen_modal.py
import modal
# import random
# from datetime import datetime
# import random
# import io
# Define the Modal image
print("Building custom image...")
image = (
modal.Image.debian_slim(python_version="3.11") # Base image
.pip_install(
"numpy",
"pandas",
"diffusers",
"transformers",
"torch",
"accelerate",
"gradio",
"safetensors",
"pillow",
) # Install Python packages
.run_commands("echo 'Image build complete!'") # Run a shell command
)
# Create a Modal app
app = modal.App("ctb-ai-img-gen-modal", image=image)
volume = modal.Volume.from_name("flux-model-vol") # Reference your volume
@app.local_entrypoint()
def generate(prompt_alias, team_color, model_alias, custom_prompt, height=360, width=640, num_inference_steps=20, guidance_scale=2.0, seed=-1):
# Debug: Print a message when the function starts
print("Starting main function inside the container...")
# Import libraries and print their versions
import numpy as np
import pandas as pd
import torch
import diffusers
import transformers
import gradio as gr
from PIL import Image as PILImage
print("Hello from Modal!")
print("NumPy version:", np.__version__)
print("Pandas version:", pd.__version__)
print("PyTorch version:", torch.__version__)
print("Diffusers version:", diffusers.__version__) # Corrected: Use the library's __version__
print("Transformers version:", transformers.__version__) # Corrected: Use the library's __version__
print("Gradio version:", gr.__version__)
print("Pillow version:", PILImage.__version__)
try:
# Generate the image
image_path, message = generate_image(prompt_alias, team_color, model_alias, custom_prompt, height, width, num_inference_steps, guidance_scale, seed)
return image_path, message
except Exception as e:
return None, f"An error occurred: {e}"
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):
import torch
from diffusers import StableDiffusionPipeline
# 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"
# 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()}")
# Format the prompt
prompt = prompt.format(team_color=team_color.lower(), enemy_color=enemy_color)
# Print the formatted prompt for debugging
print("\nFormatted Prompt:")
print(prompt)
# Append custom prompt if provided
if custom_prompt and len(custom_prompt.strip()) > 0:
prompt += " " + custom_prompt.strip()
# Randomize seed if needed
if seed == -1:
seed = random.randint(0, 1000000)
# Initialize the pipeline
pipe = StableDiffusionPipeline.from_pretrained(
model_name,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16"
)
pipe.to("cpu")
# Generate the image
try:
image = pipe(
prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
width=width,
height=height,
generator=torch.Generator("cuda").manual_seed(seed)
).images[0]
# Convert PIL image to bytes
img_byte_arr = io.BytesIO()
image.save(img_byte_arr, format='PNG')
img_byte_arr = img_byte_arr.getvalue()
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 img_byte_arr, "Image generated successfully!"
except Exception as e:
return None, f"ERROR: Failed to generate image. Details: {e}"
return output_filename, "Image generated successfully!" |