(pipeline, pipeline_params, controlnet_params, tokenizer, args, rng, weight_dtype)
| 67 | |
| 68 | |
| 69 | def log_validation(pipeline, pipeline_params, controlnet_params, tokenizer, args, rng, weight_dtype): |
| 70 | logger.info("Running validation...") |
| 71 | |
| 72 | pipeline_params = pipeline_params.copy() |
| 73 | pipeline_params["controlnet"] = controlnet_params |
| 74 | |
| 75 | num_samples = jax.device_count() |
| 76 | prng_seed = jax.random.split(rng, jax.device_count()) |
| 77 | |
| 78 | if len(args.validation_image) == len(args.validation_prompt): |
| 79 | validation_images = args.validation_image |
| 80 | validation_prompts = args.validation_prompt |
| 81 | elif len(args.validation_image) == 1: |
| 82 | validation_images = args.validation_image * len(args.validation_prompt) |
| 83 | validation_prompts = args.validation_prompt |
| 84 | elif len(args.validation_prompt) == 1: |
| 85 | validation_images = args.validation_image |
| 86 | validation_prompts = args.validation_prompt * len(args.validation_image) |
| 87 | else: |
| 88 | raise ValueError( |
| 89 | "number of `args.validation_image` and `args.validation_prompt` should be checked in `parse_args`" |
| 90 | ) |
| 91 | |
| 92 | image_logs = [] |
| 93 | |
| 94 | for validation_prompt, validation_image in zip(validation_prompts, validation_images): |
| 95 | prompts = num_samples * [validation_prompt] |
| 96 | prompt_ids = pipeline.prepare_text_inputs(prompts) |
| 97 | prompt_ids = shard(prompt_ids) |
| 98 | |
| 99 | validation_image = Image.open(validation_image).convert("RGB") |
| 100 | processed_image = pipeline.prepare_image_inputs(num_samples * [validation_image]) |
| 101 | processed_image = shard(processed_image) |
| 102 | images = pipeline( |
| 103 | prompt_ids=prompt_ids, |
| 104 | image=processed_image, |
| 105 | params=pipeline_params, |
| 106 | prng_seed=prng_seed, |
| 107 | num_inference_steps=50, |
| 108 | jit=True, |
| 109 | ).images |
| 110 | |
| 111 | images = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]) |
| 112 | images = pipeline.numpy_to_pil(images) |
| 113 | |
| 114 | image_logs.append( |
| 115 | {"validation_image": validation_image, "images": images, "validation_prompt": validation_prompt} |
| 116 | ) |
| 117 | |
| 118 | if args.report_to == "wandb": |
| 119 | formatted_images = [] |
| 120 | for log in image_logs: |
| 121 | images = log["images"] |
| 122 | validation_prompt = log["validation_prompt"] |
| 123 | validation_image = log["validation_image"] |
| 124 | |
| 125 | formatted_images.append(wandb.Image(validation_image, caption="Controlnet conditioning")) |
| 126 | for image in images: |
no test coverage detected
searching dependent graphs…