minor fix in controlnet flax example

examples/controlnet/README.md

@@ -320,6 +320,12 @@ Then cd in the example folder and run
 pip install -U -r requirements_flax.txt
 ```
 
+If you want to use Weights and Bias logging, you should also install `wandb` now
+
+```bash
+pip install wandb
+```
+
 Now let's downloading two conditioning images that we will use to run validation during the training in order to track our progress
 
 ```

examples/controlnet/train_controlnet_flax.py

@@ -154,15 +154,16 @@ def log_validation(controlnet, controlnet_params, tokenizer, args, rng, weight_d
 
 def save_model_card(repo_id: str, image_logs=None, base_model=str, repo_folder=None):
     img_str = ""
-    for i, log in enumerate(image_logs):
-        images = log["images"]
-        validation_prompt = log["validation_prompt"]
-        validation_image = log["validation_image"]
-        validation_image.save(os.path.join(repo_folder, "image_control.png"))
-        img_str += f"prompt: {validation_prompt}\n"
-        images = [validation_image] + images
-        image_grid(images, 1, len(images)).save(os.path.join(repo_folder, f"images_{i}.png"))
-        img_str += f"![images_{i})](./images_{i}.png)\n"
+    if image_logs is not None:
+        for i, log in enumerate(image_logs):
+            images = log["images"]
+            validation_prompt = log["validation_prompt"]
+            validation_image = log["validation_image"]
+            validation_image.save(os.path.join(repo_folder, "image_control.png"))
+            img_str += f"prompt: {validation_prompt}\n"
+            images = [validation_image] + images
+            image_grid(images, 1, len(images)).save(os.path.join(repo_folder, f"images_{i}.png"))
+            img_str += f"![images_{i})](./images_{i}.png)\n"
 
     yaml = f"""
 ---
@@ -1021,6 +1022,8 @@ def cumul_grad_step(grad_idx, loss_grad_rng):
     if jax.process_index() == 0:
         if args.validation_prompt is not None:
             image_logs = log_validation(controlnet, state.params, tokenizer, args, validation_rng, weight_dtype)
+        else:
+            image_logs = None
 
         controlnet.save_pretrained(
             args.output_dir,