File size: 5,168 Bytes
338d95d |
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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 |
# compi_phase1b_styled_generation.py
import os
import sys
import torch
from datetime import datetime
from diffusers import StableDiffusionPipeline
from PIL import Image
# Add project root to path for imports
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
# -------- 1. SETUP --------
if torch.cuda.is_available():
device = "cuda"
print("Running on CUDA GPU.")
else:
device = "cpu"
print("Running on CPU.")
OUTPUT_DIR = os.path.join(os.path.dirname(__file__), '..', '..', "outputs")
os.makedirs(OUTPUT_DIR, exist_ok=True)
def log(msg):
now = datetime.now().strftime("[%Y-%m-%d %H:%M:%S]")
print(f"{now} {msg}")
# -------- 2. LOAD MODEL --------
MODEL_NAME = "runwayml/stable-diffusion-v1-5"
log(f"Loading model: {MODEL_NAME}")
def dummy_safety_checker(images, **kwargs):
return images, [False] * len(images)
try:
pipe = StableDiffusionPipeline.from_pretrained(
MODEL_NAME,
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
safety_checker=dummy_safety_checker,
)
except Exception as e:
log(f"Error loading model: {e}")
sys.exit(1)
pipe = pipe.to(device)
pipe.enable_attention_slicing()
log("Model loaded.")
# -------- 3. STYLE & MOOD PROMPT ENGINEERING --------
# Predefined styles and moods (add more as desired)
STYLES = [
"digital art",
"oil painting",
"watercolor",
"cyberpunk",
"impressionist",
"concept art",
"anime",
"photorealistic",
"minimalist",
"surrealism",
"pixel art",
"steampunk",
"3d render"
]
MOODS = [
"dreamy atmosphere",
"dark and moody",
"peaceful",
"vibrant and energetic",
"melancholic",
"mysterious",
"whimsical",
"serene",
"uplifting",
"dramatic lighting",
"retro"
]
def main():
"""Main function for command-line execution"""
# Input: main prompt
if len(sys.argv) > 1:
main_prompt = " ".join(sys.argv[1:])
log(f"Prompt from command line: {main_prompt}")
else:
main_prompt = input("Enter your main scene/subject (e.g., 'A forest of bioluminescent trees'): ").strip()
if not main_prompt:
log("No main prompt entered. Exiting.")
sys.exit(0)
# Style selector
print("\nChoose an art style from the list or enter your own:")
for idx, style in enumerate(STYLES, 1):
print(f" {idx}. {style}")
style_choice = input(f"Enter style number [1-{len(STYLES)}] or type your own: ").strip()
if style_choice.isdigit() and 1 <= int(style_choice) <= len(STYLES):
style = STYLES[int(style_choice)-1]
else:
style = style_choice if style_choice else STYLES[0]
log(f"Style selected: {style}")
# Mood selector
print("\nChoose a mood from the list or enter your own:")
for idx, mood in enumerate(MOODS, 1):
print(f" {idx}. {mood}")
mood_choice = input(f"Enter mood number [1-{len(MOODS)}] or type your own (or leave blank): ").strip()
if mood_choice.isdigit() and 1 <= int(mood_choice) <= len(MOODS):
mood = MOODS[int(mood_choice)-1]
elif mood_choice:
mood = mood_choice
else:
mood = ""
log(f"Mood selected: {mood if mood else '[none]'}")
# Combine all for final prompt
full_prompt = main_prompt
if style: full_prompt += f", {style}"
if mood: full_prompt += f", {mood}"
log(f"Full prompt: {full_prompt}")
# -------- 4. GENERATION PARAMETERS --------
NUM_VARIATIONS = input("How many variations to generate? (default 1): ").strip()
try:
NUM_VARIATIONS = max(1, int(NUM_VARIATIONS))
except Exception:
NUM_VARIATIONS = 1
INFERENCE_STEPS = 30
GUIDANCE_SCALE = 7.5
HEIGHT = 512
WIDTH = 512
# -------- 5. IMAGE GENERATION --------
log(f"Generating {NUM_VARIATIONS} image(s) for prompt: '{full_prompt}'")
images = []
for i in range(NUM_VARIATIONS):
seed = torch.seed() # random seed for each variation
generator = torch.manual_seed(seed) if device == "cpu" else torch.Generator(device).manual_seed(seed)
with torch.autocast(device) if device == "cuda" else torch.no_grad():
result = pipe(
full_prompt,
height=HEIGHT,
width=WIDTH,
num_inference_steps=INFERENCE_STEPS,
guidance_scale=GUIDANCE_SCALE,
generator=generator,
)
img: Image.Image = result.images[0]
# Compose filename
prompt_slug = "_".join(main_prompt.lower().split()[:5])
style_slug = style.replace(" ", "")[:10]
mood_slug = mood.replace(" ", "")[:10] if mood else "none"
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
fname = f"{prompt_slug[:25]}_{style_slug}_{mood_slug}_{timestamp}_seed{seed}_v{i+1}.png"
fpath = os.path.join(OUTPUT_DIR, fname)
img.save(fpath)
log(f"Image saved: {fpath}")
images.append(fpath)
log(f"All {NUM_VARIATIONS} images generated and saved.")
log("Phase 1.B complete.")
if __name__ == "__main__":
main()
|