import argparse import json import os import random import logging import shutil from typing import Dict, List, Optional, Tuple from PIL import Image from tqdm import tqdm from multiprocessing import Pool, cpu_count from functools import partial from scaleup_utils import ( load_jsonl, save_jsonl, load_blended_sample, get_blended_sample_dirs, compute_non_overlapping_box_xyxy, compute_total_overlap, create_layer_on_canvas, build_spatial_aware_caption, get_position_description, get_box_size, get_content_bbox, load_caption_list, get_laion_images_with_captions, get_caption_images_with_text, select_random_layers_from_samples, ) logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Default canvas size CANVAS_SIZE = 1024 def parse_args(): parser = argparse.ArgumentParser(description='Scale up PrismLayersPro-blended dataset') parser.add_argument('--blended_dir', type=str, required=True, help='Path to PrismLayersPro-blended directory') parser.add_argument('--laion_dir', type=str, required=True, help='Path to LAION aesthetic images directory') parser.add_argument('--caption_dir', type=str, required=True, help='Path to caption images directory') parser.add_argument('--caption_meta', type=str, required=True, help='Path to captions.jsonl with caption text') parser.add_argument('--output_dir', type=str, required=True, help='Output directory for scaled-up dataset') parser.add_argument('--num_samples', type=int, default=100000, help='Number of new samples to generate') parser.add_argument('--start_index', type=int, default=0, help='Starting sample index for output naming') parser.add_argument('--seed', type=int, default=42, help='Random seed') # Layer selection parameters parser.add_argument('--min_donor_samples', type=int, default=2, help='Minimum number of samples to pick layers from') parser.add_argument('--max_donor_samples', type=int, default=3, help='Maximum number of samples to pick layers from') parser.add_argument('--min_layers_per_donor', type=int, default=1, help='Minimum layers to pick from each donor sample') parser.add_argument('--max_layers_per_donor', type=int, default=2, help='Maximum layers to pick from each donor sample') parser.add_argument('--added_layer_min_size', type=float, default=0.8, help='Minimum size ratio for added layers') parser.add_argument('--added_layer_max_size', type=float, default=1.2, help='Maximum size ratio for added layers') parser.add_argument('--laion_prob', type=float, default=0.1, help='Probability of including LAION image layer') parser.add_argument('--caption_prob', type=float, default=0.2, help='Probability of including caption text image layer') parser.add_argument('--laion_min_size', type=float, default=0.2, help='Minimum size ratio for LAION layer') parser.add_argument('--laion_max_size', type=float, default=0.4, help='Maximum size ratio for LAION layer') parser.add_argument('--caption_min_size', type=float, default=1.0, help='Minimum size ratio for caption layer') parser.add_argument('--caption_max_size', type=float, default=1.2, help='Maximum size ratio for caption layer') # Base layer removal parameters parser.add_argument('--min_layers_to_remove', type=int, default=2, help='Minimum number of layers to remove from base sample') parser.add_argument('--max_layers_to_remove', type=int, default=4, help='Maximum number of layers to remove from base sample') # AlphaVAE layer parameters parser.add_argument('--alphavae_dir', type=str, default=None, help='Path to AlphaVAE_frontview images directory') parser.add_argument('--alphavae_prompts', type=str, default=None, help='Path to prompts.txt for AlphaVAE captions') parser.add_argument('--alphavae_min_layers', type=int, default=0, help='Minimum number of AlphaVAE layers to add per sample') parser.add_argument('--alphavae_max_layers', type=int, default=0, help='Maximum number of AlphaVAE layers to add per sample') parser.add_argument('--alphavae_min_size', type=float, default=0.15, help='Minimum size ratio for AlphaVAE layers') parser.add_argument('--alphavae_max_size', type=float, default=0.35, help='Maximum size ratio for AlphaVAE layers') # Multiprocessing parser.add_argument('--max_base_samples', type=int, default=None, help='Max number of base samples to use (sorted by name). ' 'E.g. 18000 to use sample_000000..sample_017999') parser.add_argument('--skip_existing', action='store_true', help='Skip samples whose output directory already exists (for resuming)') parser.add_argument('--num_workers', type=int, default=64, help='Number of parallel workers (0 = single-process)') return parser.parse_args() def create_scaled_up_sample( base_sample_dir: str, all_sample_dirs: List[str], laion_images: List[Tuple[str, str]], caption_images: List[Tuple[str, str]], alphavae_images: List[Tuple[str, str]], output_dir: str, sample_idx: int, args: argparse.Namespace, ) -> Optional[Dict]: """ Create a new scaled-up sample by combining a base sample with layers from other samples. Returns metadata dict or None if failed. """ # Load base sample base_meta = load_blended_sample(base_sample_dir) if base_meta is None: logger.warning(f"Failed to load base sample: {base_sample_dir}") return None canvas_size = base_meta.get('width', CANVAS_SIZE) # Create output sample directory sample_name = f"sample_{sample_idx:06d}" sample_output_dir = os.path.join(output_dir, sample_name) os.makedirs(sample_output_dir, exist_ok=True) # Copy base_image base_image = base_meta.get('base_image') if base_image: base_image.save(os.path.join(sample_output_dir, 'base_image.png')) else: base_image = Image.new('RGBA', (canvas_size, canvas_size), (0, 0, 0, 0)) base_image.save(os.path.join(sample_output_dir, 'base_image.png')) # Start with base composite composite = base_image.copy() # Collect occupied boxes occupied_boxes = [] # New layers list new_layers = [] current_layer_idx = 0 # === Step 1: Copy layers from base sample (excluding laion_foreground and caption types) === # Also randomly remove 2-3 layers to keep total layer count reasonable base_layers = base_meta.get('layers', []) base_prism_layers = [l for l in base_layers if l.get('type') is None] # Randomly remove some layers from base num_to_remove = random.randint(args.min_layers_to_remove, args.max_layers_to_remove) num_to_remove = min(num_to_remove, max(0, len(base_prism_layers) - 1)) # Keep at least 1 layer removed_layer_indices = set() if num_to_remove > 0 and len(base_prism_layers) > 1: layers_to_remove = random.sample(base_prism_layers, num_to_remove) removed_layer_indices = {l['layer_idx'] for l in layers_to_remove} # Filter out removed layers base_prism_layers_filtered = [l for l in base_prism_layers if l['layer_idx'] not in removed_layer_indices] for layer in base_prism_layers_filtered: orig_layer_idx = layer['layer_idx'] layer_img = base_meta.get('layer_images', {}).get(orig_layer_idx) if layer_img is None: continue orig_box = layer.get('box', [0, 0, canvas_size, canvas_size]) caption = layer.get('caption', '') orig_w = orig_box[2] - orig_box[0] orig_h = orig_box[3] - orig_box[1] # Randomly reposition (layout-agnostic): find a non-overlapping spot best_box = None best_overlap_ratio = float('inf') new_box = None for _ in range(300): x0 = random.randint(0, max(0, canvas_size - orig_w)) y0 = random.randint(0, max(0, canvas_size - orig_h)) candidate = [x0, y0, x0 + orig_w, y0 + orig_h] box_area = orig_w * orig_h if box_area <= 0: continue overlap = compute_total_overlap(candidate, occupied_boxes) overlap_ratio = overlap / box_area if overlap == 0: new_box = candidate break if overlap_ratio < best_overlap_ratio: best_overlap_ratio = overlap_ratio best_box = candidate if new_box is None: new_box = best_box if best_box else [0, 0, orig_w, orig_h] # Place cropped layer onto full canvas at the new random position layer_canvas = create_layer_on_canvas(layer_img, new_box, canvas_size) # Save layer with new index layer_filename = f'layer_{current_layer_idx:02d}.png' layer_canvas.save(os.path.join(sample_output_dir, layer_filename)) # Composite composite = Image.alpha_composite(composite, layer_canvas) # Record w, h = get_box_size(new_box) new_layers.append({ 'layer_idx': current_layer_idx, 'caption': caption, 'box': new_box, 'width_dst': w, 'height_dst': h, 'image_path': layer_filename, 'source': 'base', 'source_sample': base_sample_dir, }) occupied_boxes.append(new_box) current_layer_idx += 1 # === Step 2: Add layers from other samples === num_donors = random.randint(args.min_donor_samples, args.max_donor_samples) donor_layers = select_random_layers_from_samples( all_sample_dirs, exclude_sample=base_sample_dir, num_samples_to_pick=num_donors, num_layers_per_sample=(args.min_layers_per_donor, args.max_layers_per_donor) ) for layer_img, layer_info, source_sample in donor_layers: caption = layer_info.get('caption', '') # Use the donor layer's original bounding box dimensions orig_box = layer_info.get('box', [0, 0, canvas_size, canvas_size]) orig_w = orig_box[2] - orig_box[0] orig_h = orig_box[3] - orig_box[1] # Find a non-overlapping position for a box of this exact size best_box = None best_overlap_ratio = float('inf') new_box = None for _ in range(300): x0 = random.randint(0, max(0, canvas_size - orig_w)) y0 = random.randint(0, max(0, canvas_size - orig_h)) candidate = [x0, y0, x0 + orig_w, y0 + orig_h] box_area = orig_w * orig_h if box_area <= 0: continue overlap = compute_total_overlap(candidate, occupied_boxes) overlap_ratio = overlap / box_area if overlap == 0: new_box = candidate break if overlap_ratio < best_overlap_ratio: best_overlap_ratio = overlap_ratio best_box = candidate if new_box is None: new_box = best_box if best_box else [0, 0, orig_w, orig_h] # Create layer on canvas at new position layer_canvas = create_layer_on_canvas(layer_img, new_box, canvas_size) # Save layer layer_filename = f'layer_{current_layer_idx:02d}.png' layer_canvas.save(os.path.join(sample_output_dir, layer_filename)) # Composite composite = Image.alpha_composite(composite, layer_canvas) # Record w, h = get_box_size(new_box) new_layers.append({ 'layer_idx': current_layer_idx, 'caption': caption, 'box': new_box, 'width_dst': w, 'height_dst': h, 'image_path': layer_filename, 'source': 'donor', 'source_sample': source_sample, 'original_layer_idx': layer_info.get('layer_idx'), }) occupied_boxes.append(new_box) current_layer_idx += 1 # === Step 2.5: Optionally add AlphaVAE layers (0 to max) === if alphavae_images and args.alphavae_max_layers > 0: num_alpha = random.randint(args.alphavae_min_layers, args.alphavae_max_layers) if num_alpha > 0: selected_alpha = random.sample(alphavae_images, min(num_alpha, len(alphavae_images))) for alpha_path, alpha_caption in selected_alpha: try: alpha_img = Image.open(alpha_path).convert('RGBA') except Exception as e: logger.warning(f"Failed to load AlphaVAE image: {alpha_path}, {e}") continue alpha_box = compute_non_overlapping_box_xyxy( canvas_size, occupied_boxes, min_size_ratio=args.alphavae_min_size, max_size_ratio=args.alphavae_max_size, max_attempts=300, max_overlap_ratio=0.10, center_margin=32 ) alpha_layer = create_layer_on_canvas(alpha_img, alpha_box, canvas_size) layer_filename = f'layer_{current_layer_idx:02d}.png' alpha_layer.save(os.path.join(sample_output_dir, layer_filename)) composite = Image.alpha_composite(composite, alpha_layer) w, h = get_box_size(alpha_box) new_layers.append({ 'layer_idx': current_layer_idx, 'caption': alpha_caption, 'box': alpha_box, 'width_dst': w, 'height_dst': h, 'image_path': layer_filename, 'type': 'alphavae', 'source_path': alpha_path, }) occupied_boxes.append(alpha_box) current_layer_idx += 1 # === Step 3: Optionally add LAION image === laion_caption = None laion_path = None if random.random() < args.laion_prob and laion_images: laion_path, laion_caption = random.choice(laion_images) try: laion_img = Image.open(laion_path).convert('RGBA') laion_orig_size = laion_img.size except Exception as e: logger.warning(f"Failed to load LAION image: {laion_path}, {e}") laion_img = None if laion_img is not None: laion_box = compute_non_overlapping_box_xyxy( canvas_size, occupied_boxes, min_size_ratio=args.laion_min_size, max_size_ratio=args.laion_max_size, max_attempts=300, max_overlap_ratio=0.10, center_margin=32 ) # Create layer laion_layer = create_layer_on_canvas(laion_img, laion_box, canvas_size) # Save layer_filename = f'layer_{current_layer_idx:02d}.png' laion_layer.save(os.path.join(sample_output_dir, layer_filename)) # Composite composite = Image.alpha_composite(composite, laion_layer) # Record w, h = get_box_size(laion_box) new_layers.append({ 'layer_idx': current_layer_idx, 'caption': laion_caption, 'box': laion_box, 'width_dst': w, 'height_dst': h, 'image_path': layer_filename, 'type': 'laion_foreground', 'original_size': list(laion_orig_size), }) occupied_boxes.append(laion_box) current_layer_idx += 1 # === Step 4: Optionally add caption image === caption_text = None caption_path = None if random.random() < args.caption_prob and caption_images: caption_path, caption_text = random.choice(caption_images) try: caption_img = Image.open(caption_path).convert('RGBA') caption_orig_size = caption_img.size except Exception as e: logger.warning(f"Failed to load caption image: {caption_path}, {e}") caption_img = None if caption_img is not None: caption_box = compute_non_overlapping_box_xyxy( canvas_size, occupied_boxes, min_size_ratio=args.caption_min_size, max_size_ratio=args.caption_max_size, max_attempts=300, max_overlap_ratio=0.10, center_margin=32 ) # Create layer caption_layer = create_layer_on_canvas(caption_img, caption_box, canvas_size) # Compute tight bbox from actual non-transparent content # Caption images have rectangular colored areas surrounded by transparency, # so we use the actual content bounds instead of the placement box. tight_box = get_content_bbox(caption_layer) if tight_box is not None: caption_box = tight_box # Save layer_filename = f'layer_{current_layer_idx:02d}.png' caption_layer.save(os.path.join(sample_output_dir, layer_filename)) # Composite composite = Image.alpha_composite(composite, caption_layer) # Record w, h = get_box_size(caption_box) new_layers.append({ 'layer_idx': current_layer_idx, 'caption': f"Text: {caption_text}" if caption_text else "Text", 'box': caption_box, 'width_dst': w, 'height_dst': h, 'image_path': layer_filename, 'type': 'caption', 'original_size': list(caption_orig_size), }) occupied_boxes.append(caption_box) current_layer_idx += 1 # === Step 5: Save whole_image === composite.save(os.path.join(sample_output_dir, 'whole_image.png')) # === Step 6: Build spatial-aware whole caption === base_caption = base_meta.get('base_caption', '') whole_caption = build_spatial_aware_caption(new_layers, canvas_size, base_caption) # === Step 7: Create metadata === metadata = { 'id': f'{sample_idx:09d}', 'style_category': base_meta.get('style_category', ''), 'whole_caption': whole_caption, 'base_caption': base_caption, 'layer_count': len(new_layers), 'layers': new_layers, # Extra fields 'sample_dir': sample_name, 'width': canvas_size, 'height': canvas_size, 'base_sample': base_sample_dir, 'num_base_layers_removed': len(removed_layer_indices), 'num_donor_samples': num_donors, 'num_donated_layers': len(donor_layers), } if laion_path: metadata['laion_path'] = laion_path metadata['laion_caption'] = laion_caption if caption_path: metadata['caption_path'] = caption_path metadata['caption_text'] = caption_text # Save metadata with open(os.path.join(sample_output_dir, 'metadata.json'), 'w', encoding='utf-8') as f: json.dump(metadata, f, indent=2, ensure_ascii=False) return metadata def _worker_fn(task): """Worker function for multiprocessing. Each task is a dict with all needed info.""" sample_idx = task['sample_idx'] base_sample_dir = task['base_sample_dir'] all_sample_dirs = task['all_sample_dirs'] laion_images = task['laion_images'] caption_images = task['caption_images'] alphavae_images = task['alphavae_images'] output_dir = task['output_dir'] args = task['args'] seed = task['seed'] if args.skip_existing: meta_path = os.path.join(output_dir, f"sample_{sample_idx:06d}", 'metadata.json') if os.path.exists(meta_path): try: with open(meta_path, 'r') as f: return json.load(f) except Exception: pass random.seed(seed) try: metadata = create_scaled_up_sample( base_sample_dir=base_sample_dir, all_sample_dirs=all_sample_dirs, laion_images=laion_images, caption_images=caption_images, alphavae_images=alphavae_images, output_dir=output_dir, sample_idx=sample_idx, args=args, ) return metadata except Exception as e: logger.error(f"Error processing sample {sample_idx}: {e}") return None def main(): args = parse_args() random.seed(args.seed) os.makedirs(args.output_dir, exist_ok=True) # Load existing blended samples logger.info("Loading existing blended samples...") all_sample_dirs = get_blended_sample_dirs(args.blended_dir, max_samples=args.max_base_samples) logger.info(f"Found {len(all_sample_dirs)} existing samples" + (f" (limited to first {args.max_base_samples})" if args.max_base_samples else "")) if len(all_sample_dirs) < 10: logger.error("Not enough existing samples to create scaled-up dataset!") return # Load caption list logger.info("Loading caption list from captions.jsonl...") caption_list = load_caption_list(args.caption_meta) logger.info(f"Loaded {len(caption_list)} caption entries") # Load LAION images (cap at 20000 for balanced diversity) logger.info("Loading LAION images with captions...") laion_images = get_laion_images_with_captions(args.laion_dir) if len(laion_images) > 20000: random.shuffle(laion_images) laion_images = laion_images[:20000] logger.info(f"Using {len(laion_images)} LAION images") # Load caption images logger.info("Loading caption images...") caption_images = get_caption_images_with_text(args.caption_dir, caption_list) logger.info(f"Found {len(caption_images)} caption images") # Load AlphaVAE images (optional) alphavae_images = [] # list of (image_path, caption) if args.alphavae_dir and args.alphavae_prompts: logger.info("Loading AlphaVAE images with prompts...") with open(args.alphavae_prompts, 'r') as f: alphavae_prompts = [l.strip() for l in f.readlines() if l.strip()] alpha_files = sorted([ f for f in os.listdir(args.alphavae_dir) if f.endswith('.png') ]) for fname in alpha_files: idx = int(fname.replace('.png', '')) prompt_idx = idx // 5 if prompt_idx < len(alphavae_prompts): caption = alphavae_prompts[prompt_idx] else: caption = "" alphavae_images.append((os.path.join(args.alphavae_dir, fname), caption)) logger.info(f"Found {len(alphavae_images)} AlphaVAE images") # Pre-generate tasks with deterministic per-sample seeds rng = random.Random(args.seed) tasks = [] for i in range(args.num_samples): sample_idx = args.start_index + i tasks.append({ 'sample_idx': sample_idx, 'base_sample_dir': rng.choice(all_sample_dirs), 'all_sample_dirs': all_sample_dirs, 'laion_images': laion_images, 'caption_images': caption_images, 'alphavae_images': alphavae_images, 'output_dir': args.output_dir, 'args': args, 'seed': rng.randint(0, 2**31), }) # Generate samples all_metadata = [] failed_count = 0 num_workers = args.num_workers if num_workers > 0: logger.info(f"Using multiprocessing with {num_workers} workers") with Pool(processes=num_workers) as pool: for metadata in tqdm( pool.imap_unordered(_worker_fn, tasks), total=len(tasks), desc="Generating samples" ): if metadata: all_metadata.append(metadata) else: failed_count += 1 else: logger.info("Using single-process mode") for task in tqdm(tasks, desc="Generating samples"): metadata = _worker_fn(task) if metadata: all_metadata.append(metadata) else: failed_count += 1 # Sort by sample index for deterministic output order all_metadata.sort(key=lambda m: int(m['id'])) # Save index (scaleup_meta.jsonl) index_path = os.path.join(args.output_dir, 'scaleup_meta.jsonl') with open(index_path, 'w', encoding='utf-8') as f: for meta in all_metadata: f.write(json.dumps(meta, ensure_ascii=False) + '\n') logger.info(f"Generated {len(all_metadata)} samples ({failed_count} failed)") logger.info(f"Output saved to {args.output_dir}") logger.info(f"Index saved to {index_path}") if __name__ == '__main__': main()