| """Visualize the diffusion generation process - capture images at each timestep.""" |
| import torch |
| import argparse |
| import os |
| import matplotlib.pyplot as plt |
|
|
| import config |
| from model import TextConditionedUNet |
| from scheduler import SimpleDDPMScheduler |
| from text_encoder import CLIPTextEncoder |
| from generate import tensor_to_image |
|
|
|
|
| def sample_with_snapshots(scheduler, model, shape, text_embeddings, device='cuda', |
| guidance_scale=1.0, snapshot_steps=None): |
| """Modified sampling that captures snapshots at specific timesteps.""" |
| b = shape[0] |
| img = torch.randn(shape, device=device) |
|
|
| |
| if snapshot_steps is None: |
| interval = scheduler.num_timesteps // 10 |
| snapshot_steps = list(range(scheduler.num_timesteps - 1, -1, -interval)) |
| if 0 not in snapshot_steps: |
| snapshot_steps.append(0) |
|
|
| snapshots = {} |
|
|
| for i in reversed(range(0, scheduler.num_timesteps)): |
| t = torch.full((b,), i, device=device, dtype=torch.long) |
| img = scheduler.p_sample_text(model, img, t, text_embeddings, guidance_scale) |
| img = torch.clamp(img, -2.0, 2.0) |
|
|
| if i in snapshot_steps: |
| snapshots[i] = img.clone().detach() |
|
|
| return img, snapshots |
|
|
|
|
| def plot_denoising_process(snapshots, prompt, output_path, sample_idx=0): |
| """Plot snapshots side by side showing noise -> final image.""" |
| timesteps = sorted(snapshots.keys(), reverse=True) |
| num_steps = len(timesteps) |
|
|
| fig, axes = plt.subplots(1, num_steps, figsize=(2.5 * num_steps, 3)) |
| if num_steps == 1: |
| axes = [axes] |
|
|
| fig.suptitle(f'Denoising Process: "{prompt}"', fontsize=12, fontweight='bold') |
|
|
| for idx, t in enumerate(timesteps): |
| img_tensor = snapshots[t][sample_idx] |
| img = tensor_to_image(img_tensor) |
|
|
| axes[idx].imshow(img, cmap='gray') |
| axes[idx].axis('off') |
| axes[idx].set_title(f't={t}' if t > 0 else 'Final', fontsize=10) |
|
|
| plt.tight_layout() |
| plt.savefig(output_path, dpi=150, bbox_inches='tight') |
| plt.close() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='Visualize denoising process') |
| parser.add_argument('--checkpoint', type=str, required=True) |
| parser.add_argument('--prompt', type=str, default="a drawing of a cat") |
| parser.add_argument('--guidance-scale', type=float, default=config.CFG_GUIDANCE_SCALE) |
| parser.add_argument('--num-steps', type=int, default=10, |
| help='Number of snapshots to capture') |
| parser.add_argument('--device', type=str, default='cuda') |
| args = parser.parse_args() |
|
|
| if args.device == 'cuda' and not torch.cuda.is_available(): |
| args.device = 'cpu' |
|
|
| |
| checkpoint = torch.load(args.checkpoint, map_location=args.device) |
| ckpt_config = checkpoint.get('config', {}) |
|
|
| model = TextConditionedUNet(text_dim=ckpt_config.get('text_dim', config.TEXT_DIM)).to(args.device) |
| model.load_state_dict(checkpoint['model_state_dict']) |
| model.eval() |
|
|
| text_encoder = CLIPTextEncoder( |
| model_name=ckpt_config.get('clip_model', config.CLIP_MODEL), freeze=True |
| ).to(args.device) |
| text_encoder.eval() |
|
|
| scheduler = SimpleDDPMScheduler(config.TIMESTEPS) |
|
|
| |
| with torch.no_grad(): |
| text_embedding = text_encoder(args.prompt) |
| shape = (1, 1, config.IMAGE_SIZE, config.IMAGE_SIZE) |
|
|
| _, snapshots = sample_with_snapshots( |
| scheduler, model, shape, text_embedding, args.device, args.guidance_scale |
| ) |
|
|
| |
| os.makedirs("outputs", exist_ok=True) |
| safe_prompt = "".join(c if c.isalnum() or c in " _" else "" for c in args.prompt)[:50] |
| output_path = f"outputs/denoising_{safe_prompt}.png" |
|
|
| plot_denoising_process(snapshots, args.prompt, output_path) |
| print(f"✅ Saved visualization: {output_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
