hidden_characters.py

"""
Hidden Characters - Diffusion Illusions
Converted from hidden_characters_for_colab.ipynb to Python script

This script generates hidden character illusions where multiple images (a, b, c, d)
overlay to reveal a hidden image (z).

Original notebook: https://github.com/RyannDaGreat/Diffusion-Illusions
"""

import numpy as np
import rp
import torch
import torch.nn as nn
import source.stable_diffusion as sd
from easydict import EasyDict
from source.learnable_textures import LearnableImageFourier
from source.stable_diffusion_labels import NegativeLabel
from itertools import chain
import time
import os
from pathlib import Path


def simulate_overlay(a, b, c, d, clean_mode=True):
    """Simulate the overlay of multiple images."""
    if clean_mode:
        exp = 1
        brightness = 3
        black = 0
    else:
        exp = rp.random_float(.5, 1)
        brightness = rp.random_float(1, 5)
        black = rp.random_float(0, .5)
        # Note: bottom and top are not defined in clean mode, so this is commented out
        # bottom = rp.blend(bottom, black, rp.random_float())
        # top = rp.blend(top, black, rp.random_float())
    return (a**exp * b**exp * c**exp * d**exp * brightness).clamp(0, 99).tanh()


def main():
    """Main function to run hidden characters generation."""
    
    # Configuration
    print("=" * 60)
    print("Hidden Characters - Diffusion Illusions")
    print("=" * 60)
    print()
    
    # Set prompts
    print("Loading example prompts...")
    example_prompts = rp.load_yaml_file('source/example_prompts.yaml')
    print('Available example prompts:', ', '.join(example_prompts))
    print()
    
    # These prompts are all strings - you can replace them with whatever you want!
    # By default it lets you choose from example prompts
    prompt_a, prompt_b, prompt_c, prompt_d, prompt_z = rp.gather(
        example_prompts, 
        'miku froggo lipstick kitten_in_box darth_vader'.split()
    )
    # Prompts a,b,c,d are the normal looking images
    # Prompt z is the hidden image you get when you overlay them all on top of each other
    
    negative_prompt = ''
    
    print()
    print('Negative prompt:', repr(negative_prompt))
    print()
    print('Chosen prompts:')
    print('    prompt_a =', repr(prompt_a))
    print('    prompt_b =', repr(prompt_b))
    print('    prompt_c =', repr(prompt_c))
    print('    prompt_d =', repr(prompt_d))
    print('    prompt_z =', repr(prompt_z))
    print()
    
    # Load Stable Diffusion model
    print("Loading Stable Diffusion model...")
    print("(This may take a few minutes on first run to download the model)")
    model_name = "CompVis/stable-diffusion-v1-4"
    
    # GPU configuration
    # Set to 'cuda:0' for first GPU, 'cuda:1' for second GPU, etc.
    # Use 'cpu' if you want to run on CPU (very slow!)
    gpu = 'cuda:0'  # Using GPU1: NVIDIA GeForce RTX 4060 laptop GPU
    
    # Check if CUDA is available
    if not torch.cuda.is_available():
        print("WARNING: CUDA not available, using CPU (will be very slow!)")
        gpu = 'cpu'
    else:
        # Extract GPU index from 'cuda:X' format
        if gpu.startswith('cuda:'):
            gpu_index = int(gpu.split(':')[1])
            gpu_name = torch.cuda.get_device_name(gpu_index)
            print(f"Using GPU{gpu_index}: {gpu_name}")
            print(f"Total CUDA devices available: {torch.cuda.device_count()}")
            
            # Verify the GPU is accessible
            try:
                test_tensor = torch.tensor([1.0]).to(gpu)
                print(f"GPU{gpu_index} is accessible and ready!")
            except Exception as e:
                print(f"WARNING: Cannot access GPU{gpu_index}, error: {e}")
                print("Falling back to CPU...")
                gpu = 'cpu'
        else:
            print(f"Using device: {gpu}")
    
    s = sd.StableDiffusion(gpu, model_name)
    device = s.device
    print("Model loaded successfully!")
    print()
    
    # Create labels
    print("Creating labels...")
    label_a = NegativeLabel(prompt_a, negative_prompt)
    label_b = NegativeLabel(prompt_b, negative_prompt)
    label_c = NegativeLabel(prompt_c, negative_prompt)
    label_d = NegativeLabel(prompt_d, negative_prompt)
    label_z = NegativeLabel(prompt_z, negative_prompt)
    print("Labels created!")
    print()
    
    # Image Parametrization and Initialization
    print("Initializing learnable images...")
    print("(This section uses VRAM)")
    
    # Select Learnable Image Size (this has big VRAM implications!):
    # Note: We use implicit neural representations for better image quality
    # They're previously used in our paper "TRITON: Neural Neural Textures make Sim2Real Consistent" 
    # ... and that representation is based on Fourier Feature Networks
    learnable_image_maker = lambda: LearnableImageFourier(
        height=256, width=256, hidden_dim=256, num_features=128
    ).to(s.device)
    SIZE = 256
    # For higher quality (but requires more VRAM):
    # learnable_image_maker = lambda: LearnableImageFourier(
    #     height=512, width=512, num_features=256, hidden_dim=256, scale=20
    # ).to(s.device)
    # SIZE = 512
    
    image_a = learnable_image_maker()
    image_b = learnable_image_maker()
    image_c = learnable_image_maker()
    image_d = learnable_image_maker()
    print(f"Images initialized (size: {SIZE}x{SIZE})!")
    print()
    
    # Setup overlay simulation
    CLEAN_MODE = True  # If False, we augment the images by randomly simulating printer quality
    
    learnable_image_a = lambda: image_a()
    learnable_image_b = lambda: image_b()
    learnable_image_c = lambda: image_c()
    learnable_image_d = lambda: image_d()
    learnable_image_z = lambda: simulate_overlay(
        image_a(), image_b(), image_c(), image_d(), clean_mode=CLEAN_MODE
    )
    
    params = chain(
        image_a.parameters(),
        image_b.parameters(),
        image_c.parameters(),
        image_d.parameters(),
    )
    optim = torch.optim.SGD(params, lr=1e-4)
    
    # Setup labels and weights
    labels = [label_a, label_b, label_c, label_d, label_z]
    learnable_images = [
        learnable_image_a, 
        learnable_image_b, 
        learnable_image_c, 
        learnable_image_d, 
        learnable_image_z
    ]
    
    # The weight coefficients for each prompt. 
    # For example, if we have [1,1,1,1,5], then the hidden prompt (prompt_z) will be prioritized
    weights = [1, 1, 1, 1, 1]
    
    weights = rp.as_numpy_array(weights)
    weights = weights / weights.sum()
    weights = weights * len(weights)
    
    # For saving a timelapse
    ims = []
    
    def get_display_image():
        """Get a tiled display of all images."""
        return rp.tiled_images(
            [
                *[rp.as_numpy_image(image()) for image in learnable_images[:-1]],
                rp.as_numpy_image(learnable_image_z()),
            ],
            length=len(learnable_images),
            border_thickness=0,
        )
    
    # Training configuration
    NUM_ITER = 10000
    
    # Set the minimum and maximum noise timesteps for the dream loss (aka score distillation loss)
    s.max_step = MAX_STEP = 990
    s.min_step = MIN_STEP = 10
    
    display_eta = rp.eta(NUM_ITER, title='Status: ')
    
    DISPLAY_INTERVAL = 200
    
    print('=' * 60)
    print('Starting training...')
    print(f'Every {DISPLAY_INTERVAL} iterations we display an image in the form [image_a, image_b, image_c, image_d, image_z] where')
    print('    image_z = image_a * image_b * image_c * image_d')
    print()
    print('Press Ctrl+C at any time to stop early and save results')
    print('You can resume training later by running this script again')
    print()
    print(f'Total iterations: {NUM_ITER}')
    print('Please expect this to take hours to get good images. The longer you wait the better they\'ll be.')
    print('=' * 60)
    print()
    
    # Training loop
    try:
        for iter_num in range(NUM_ITER):
            display_eta(iter_num)  # Print the remaining time
            
            preds = []
            for label, learnable_image, weight in rp.random_batch(
                list(zip(labels, learnable_images, weights)), 
                batch_size=1
            ):
                pred = s.train_step(
                    label.embedding,
                    learnable_image()[None],
                    
                    # PRESETS (uncomment one):
                    noise_coef=.1 * weight, guidance_scale=60,  # Default
                    # noise_coef=0, image_coef=-.01, guidance_scale=50,
                    # noise_coef=0, image_coef=-.005, guidance_scale=50,
                    # noise_coef=.1, image_coef=-.010, guidance_scale=50,
                    # noise_coef=.1, image_coef=-.005, guidance_scale=50,
                    # noise_coef=.1 * weight, image_coef=-.005 * weight, guidance_scale=50,
                )
                preds += list(pred)
            
            # Save and display progress
            if not iter_num % DISPLAY_INTERVAL:
                im = get_display_image()
                ims.append(im)
                
                # Display progress
                progress = f"[{iter_num}/{NUM_ITER}] ({iter_num*100/NUM_ITER:.1f}%)"
                print(f"\n{progress} - Displaying current state...")
                
                # Save intermediate result
                output_dir = Path("outputs/hidden_characters")
                output_dir.mkdir(parents=True, exist_ok=True)
                intermediate_path = output_dir / f"progress_{iter_num:05d}.png"
                rp.save_image(im, str(intermediate_path))
                print(f"Saved: {intermediate_path}")
            
            optim.step()
            optim.zero_grad()
            
    except KeyboardInterrupt:
        print()
        print(f'Interrupted early at iteration {iter_num}')
        im = get_display_image()
        ims.append(im)
    
    # Save final results
    print()
    print('=' * 60)
    print('Training complete! Saving final results...')
    print('=' * 60)
    
    output_dir = Path("outputs/hidden_characters")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Save individual images
    print()
    print('Image A')
    img_a = rp.as_numpy_image(learnable_image_a())
    rp.save_image(img_a, str(output_dir / "image_a.png"))
    print(f"Saved: {output_dir / 'image_a.png'}")
    
    print('Image B')
    img_b = rp.as_numpy_image(learnable_image_b())
    rp.save_image(img_b, str(output_dir / "image_b.png"))
    print(f"Saved: {output_dir / 'image_b.png'}")
    
    print('Image C')
    img_c = rp.as_numpy_image(learnable_image_c())
    rp.save_image(img_c, str(output_dir / "image_c.png"))
    print(f"Saved: {output_dir / 'image_c.png'}")
    
    print('Image D')
    img_d = rp.as_numpy_image(learnable_image_d())
    rp.save_image(img_d, str(output_dir / "image_d.png"))
    print(f"Saved: {output_dir / 'image_d.png'}")
    
    print('Image Z (overlay result)')
    img_z = rp.as_numpy_image(learnable_image_z())
    rp.save_image(img_z, str(output_dir / "image_z.png"))
    print(f"Saved: {output_dir / 'image_z.png'}")
    
    # Save timelapse if we have multiple images
    if len(ims) > 1:
        timelapse_dir = output_dir / f"timelapse_{int(time.time())}"
        timelapse_dir.mkdir(parents=True, exist_ok=True)
        
        ims_names = [f'ims_{i:04d}.png' for i in range(len(ims))]
        with rp.SetCurrentDirectoryTemporarily(timelapse_dir):
            rp.save_images(ims, ims_names, show_progress=True)
        print()
        print(f'Saved timelapse to folder: {timelapse_dir}')
    
    # Save final tiled image
    final_tiled = get_display_image()
    rp.save_image(final_tiled, str(output_dir / "final_tiled.png"))
    print(f"Saved final tiled image: {output_dir / 'final_tiled.png'}")
    
    print()
    print('=' * 60)
    print('All results saved to:', output_dir.absolute())
    print('=' * 60)


if __name__ == "__main__":
    main()