# compi_phase1b_advanced_styling.py # Advanced style conditioning with negative prompts, quality settings, and enhanced prompt engineering import os import sys import torch import json import argparse 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 AND ARGUMENT PARSING -------- def setup_args(): parser = argparse.ArgumentParser(description="CompI Phase 1.B: Advanced Style Conditioning") parser.add_argument("prompt", nargs="*", help="Main scene/subject description") parser.add_argument("--style", "-s", help="Art style (or number from list)") parser.add_argument("--mood", "-m", help="Mood/atmosphere (or number from list)") parser.add_argument("--variations", "-v", type=int, default=1, help="Number of variations") parser.add_argument("--quality", "-q", choices=["draft", "standard", "high"], default="standard", help="Quality preset") parser.add_argument("--negative", "-n", help="Negative prompt") parser.add_argument("--interactive", "-i", action="store_true", help="Interactive mode") parser.add_argument("--list-styles", action="store_true", help="List available styles and exit") parser.add_argument("--list-moods", action="store_true", help="List available moods and exit") return parser.parse_args() # Device setup device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") 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. STYLE AND MOOD DEFINITIONS -------- STYLES = { "digital art": { "prompt": "digital art, highly detailed", "negative": "blurry, pixelated, low resolution" }, "oil painting": { "prompt": "oil painting, classical art, brushstrokes, canvas texture", "negative": "digital, pixelated, modern" }, "watercolor": { "prompt": "watercolor painting, soft colors, flowing paint", "negative": "harsh lines, digital, photographic" }, "cyberpunk": { "prompt": "cyberpunk style, neon lights, futuristic, sci-fi", "negative": "natural, organic, pastoral" }, "impressionist": { "prompt": "impressionist painting, soft brushstrokes, light and color", "negative": "sharp details, photorealistic, digital" }, "concept art": { "prompt": "concept art, professional illustration, detailed", "negative": "amateur, sketch, unfinished" }, "anime": { "prompt": "anime style, manga, Japanese animation", "negative": "realistic, western cartoon, photographic" }, "photorealistic": { "prompt": "photorealistic, high detail, professional photography", "negative": "cartoon, painting, stylized" }, "minimalist": { "prompt": "minimalist art, clean lines, simple composition", "negative": "cluttered, complex, detailed" }, "surrealism": { "prompt": "surrealist art, dreamlike, impossible, Salvador Dali style", "negative": "realistic, logical, mundane" }, "pixel art": { "prompt": "pixel art, 8-bit style, retro gaming", "negative": "smooth, high resolution, photorealistic" }, "steampunk": { "prompt": "steampunk style, Victorian era, brass and copper, gears", "negative": "modern, digital, futuristic" }, "3d render": { "prompt": "3D render, CGI, computer graphics, ray tracing", "negative": "2D, flat, hand-drawn" } } MOODS = { "dreamy": { "prompt": "dreamy atmosphere, soft lighting, ethereal", "negative": "harsh, stark, realistic" }, "dark": { "prompt": "dark and moody, dramatic shadows, mysterious", "negative": "bright, cheerful, light" }, "peaceful": { "prompt": "peaceful, serene, calm, tranquil", "negative": "chaotic, violent, disturbing" }, "vibrant": { "prompt": "vibrant and energetic, bright colors, dynamic", "negative": "dull, muted, lifeless" }, "melancholic": { "prompt": "melancholic, sad, nostalgic, wistful", "negative": "happy, joyful, upbeat" }, "mysterious": { "prompt": "mysterious, enigmatic, hidden secrets", "negative": "obvious, clear, straightforward" }, "whimsical": { "prompt": "whimsical, playful, fantastical, magical", "negative": "serious, realistic, mundane" }, "dramatic": { "prompt": "dramatic lighting, high contrast, cinematic", "negative": "flat lighting, low contrast, amateur" }, "retro": { "prompt": "retro style, vintage, nostalgic, classic", "negative": "modern, contemporary, futuristic" } } QUALITY_PRESETS = { "draft": {"steps": 20, "guidance": 6.0, "size": (512, 512)}, "standard": {"steps": 30, "guidance": 7.5, "size": (512, 512)}, "high": {"steps": 50, "guidance": 8.5, "size": (768, 768)} } # -------- 3. MODEL LOADING -------- def 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, ) pipe = pipe.to(device) pipe.enable_attention_slicing() log("Model loaded successfully") return pipe except Exception as e: log(f"Error loading model: {e}") sys.exit(1) # -------- 4. INTERACTIVE FUNCTIONS -------- def list_options(options_dict, title): print(f"\n{title}:") for idx, (key, value) in enumerate(options_dict.items(), 1): prompt_preview = value["prompt"][:50] + "..." if len(value["prompt"]) > 50 else value["prompt"] print(f" {idx:2d}. {key:15s} - {prompt_preview}") def get_user_choice(options_dict, prompt_text, allow_custom=True): choice = input(f"{prompt_text}: ").strip() if choice.isdigit(): idx = int(choice) - 1 keys = list(options_dict.keys()) if 0 <= idx < len(keys): return keys[idx] if choice in options_dict: return choice if allow_custom and choice: return choice return None def interactive_mode(pipe): log("Starting interactive style conditioning mode") # Get main prompt main_prompt = input("\nEnter your main scene/subject: ").strip() if not main_prompt: log("No prompt provided") return # Show and select style list_options(STYLES, "Available Styles") style_key = get_user_choice(STYLES, "Choose style (number/name/custom)") # Show and select mood list_options(MOODS, "Available Moods") mood_key = get_user_choice(MOODS, "Choose mood (number/name/custom/blank)", allow_custom=True) # Get additional parameters variations = input("Number of variations (default 1): ").strip() variations = int(variations) if variations.isdigit() else 1 quality = input("Quality [draft/standard/high] (default standard): ").strip() quality = quality if quality in QUALITY_PRESETS else "standard" negative = input("Negative prompt (optional): ").strip() # Generate images generate_styled_images(pipe, main_prompt, style_key, mood_key, variations, quality, negative) # -------- 5. GENERATION FUNCTION -------- def generate_styled_images(pipe, main_prompt, style_key, mood_key, variations, quality, custom_negative=""): # Build the full prompt full_prompt = main_prompt style_negative = "" mood_negative = "" if style_key and style_key in STYLES: full_prompt += f", {STYLES[style_key]['prompt']}" style_negative = STYLES[style_key]['negative'] elif style_key: full_prompt += f", {style_key}" if mood_key and mood_key in MOODS: full_prompt += f", {MOODS[mood_key]['prompt']}" mood_negative = MOODS[mood_key]['negative'] elif mood_key: full_prompt += f", {mood_key}" # Build negative prompt negative_parts = [part for part in [style_negative, mood_negative, custom_negative] if part] full_negative = ", ".join(negative_parts) if negative_parts else None # Get quality settings quality_settings = QUALITY_PRESETS[quality] log(f"Full prompt: {full_prompt}") log(f"Negative prompt: {full_negative or '[none]'}") log(f"Quality: {quality} ({quality_settings['steps']} steps)") log(f"Generating {variations} variation(s)") # Generate images for i in range(variations): seed = torch.seed() 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, negative_prompt=full_negative, height=quality_settings["size"][1], width=quality_settings["size"][0], num_inference_steps=quality_settings["steps"], guidance_scale=quality_settings["guidance"], generator=generator, ) img = result.images[0] # Create filename prompt_slug = "_".join(main_prompt.lower().split()[:4]) style_slug = (style_key or "nostyle").replace(" ", "")[:10] mood_slug = (mood_key or "nomood").replace(" ", "")[:10] timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"{prompt_slug[:20]}_{style_slug}_{mood_slug}_{quality}_{timestamp}_seed{seed}_v{i+1}.png" filepath = os.path.join(OUTPUT_DIR, filename) img.save(filepath) # Save metadata metadata = { "main_prompt": main_prompt, "style": style_key, "mood": mood_key, "full_prompt": full_prompt, "negative_prompt": full_negative, "quality": quality, "seed": seed, "variation": i + 1, "timestamp": datetime.now().isoformat(), "settings": quality_settings } metadata_file = filepath.replace('.png', '_metadata.json') with open(metadata_file, 'w') as f: json.dump(metadata, f, indent=2) log(f"Generated variation {i+1}: {filepath}") log(f"Phase 1.B complete - {variations} styled images generated") # -------- 6. MAIN FUNCTION -------- def main(): args = setup_args() # Handle list commands if args.list_styles: list_options(STYLES, "Available Styles") return if args.list_moods: list_options(MOODS, "Available Moods") return # Load model pipe = load_model() # Interactive mode if args.interactive: interactive_mode(pipe) return # Command line mode main_prompt = " ".join(args.prompt) if args.prompt else input("Enter main prompt: ").strip() if not main_prompt: log("No prompt provided") return generate_styled_images( pipe, main_prompt, args.style, args.mood, args.variations, args.quality, args.negative or "" ) if __name__ == "__main__": main()