| 39 | |
| 40 | |
| 41 | def log_validation(unet, args, accelerator, weight_dtype, epoch, is_final_validation=False): |
| 42 | logger.info("Running validation... ") |
| 43 | inference_ctx = contextlib.nullcontext() if is_final_validation else torch.autocast("cuda") |
| 44 | |
| 45 | if not is_final_validation: |
| 46 | unet = accelerator.unwrap_model(unet) |
| 47 | else: |
| 48 | unet = JoDiffusionModel.from_pretrained(args.output_dir, torch_dtype=weight_dtype) |
| 49 | |
| 50 | pipeline_kwargs = { |
| 51 | "pretrained_model_name_or_path": args.pretrained_model_name_or_path, |
| 52 | "unet": unet, "torch_dtype": weight_dtype, |
| 53 | } |
| 54 | if args.lightweight_label_vae: |
| 55 | from pipelines.modeling_lightweight_vae import LightweightLabelVAE |
| 56 | pipeline_kwargs["label_vae"] = LightweightLabelVAE.from_pretrained(args.pretrained_label_vae_path, torch_dtype=weight_dtype) |
| 57 | else: |
| 58 | pipeline_kwargs["label_vae"] = AutoencoderKL.from_pretrained(args.pretrained_label_vae_path, torch_dtype=weight_dtype) |
| 59 | pipeline = JoDiffusionPipeline.from_pretrained(**pipeline_kwargs) |
| 60 | pipeline = pipeline.to(accelerator.device) |
| 61 | pipeline.set_progress_bar_config(disable=True) |
| 62 | |
| 63 | if args.enable_xformers_memory_efficient_attention: |
| 64 | pipeline.enable_xformers_memory_efficient_attention() |
| 65 | |
| 66 | generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) |
| 67 | if args.dataset_name == "ade20k_semantic": |
| 68 | prompt = "a bathroom with a toilet and a shower." |
| 69 | elif args.dataset_name == "coco_semantic": |
| 70 | prompt = "a person riding a bike on a street with a car parked on the side of the road." |
| 71 | elif args.dataset_name == "voc_semantic": |
| 72 | prompt = "a person walking a dog on a leash with a car parked on the side of the road." |
| 73 | else: |
| 74 | raise ValueError(f"Unknown dataset {args.dataset_name}") |
| 75 | |
| 76 | images = [] |
| 77 | labels = [] |
| 78 | for iii in range(4): |
| 79 | with inference_ctx: |
| 80 | sample = pipeline( |
| 81 | mode="text2img" if iii < 2 else "joint", |
| 82 | prompt=prompt, |
| 83 | num_inference_steps=50, |
| 84 | generator=generator, |
| 85 | ignore_label=0, |
| 86 | ) |
| 87 | image = sample.images[0] |
| 88 | label = sample.labels[0] |
| 89 | images.append(image) |
| 90 | labels.append(label) |
| 91 | |
| 92 | tracker_key = "test" if is_final_validation else f"validation-epoch{epoch}" |
| 93 | for tracker in accelerator.trackers: |
| 94 | if tracker.name == "wandb": |
| 95 | table_data = [] |
| 96 | |
| 97 | for idx in range(len(images)): |
| 98 | table_data.append([ |