Unexpected validation image on webdataset SDXL controlnet training script

[brycegoh]

Describe the bug

I am using the webdataset sdxl controlnet training script.

I ran the script two times with 2 different learning rates in which both runs gave unexpected generated validation images:

LR = 5e-5

Step 1050

Step 1100

LR = 4e-4

Step 50

Step 100

I understand this webdataset script differs from the example script in 2 ways:

1. The webdataset script uses the EDM formulation
2. The webdataset script also takes the image cropping top left coordinates into consideration when computing the embeddings whereas the example script takes a pre-defined crops_coords_top_left_h and crops_coords_top_left_w

Therefore, is the generated validation images expected as part of the training process?

Reproduction

These are my parameters:

! accelerate launch diffusers/examples/research_projects/controlnet/train_controlnet_webdataset.py \
--pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0" \
 --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \
 --output_dir="/workspace/output" \
 --train_shards_path_or_url="/workspace/data/dataset.tar" \
 --eval_shards_path_or_url="/workspace/data/valdiation_dataset.tar" \
 --control_type="seg" \
 --dataloader_num_workers=2 \
 --max_train_samples=20210 \
 --max_eval_samples=2000 \
 --mixed_precision="bf16" \
 --resolution=1024 \
 --learning_rate=5e-5 \
 --lr_scheduler="constant" \
 --max_train_steps=33690 \
 --validation_image "/workspace/segs/image1.png" \
 --validation_prompt "a living room with a red carpet and white walls"\
 --validation_steps=50 \
 --train_batch_size=6 \
 --gradient_accumulation_steps=1 \
 --report_to="wandb" \
 --tracker_project_name="sdxl-controlnet-run-02" \
 --seed=42 \
 --hub_model_id="helloHearthXHome/cn" \
 --checkpointing_steps=1685 \

Logs

No response

System Info

Diffusers v0.29.1

Who can help?

@sayakpaul @dn

[rphly]

also facing some issues, getting black images

[sayakpaul]

That script belongs to research_projects. So, we cannot guarantee its maintenance. So, I am going to convert this issue to a discussion and provide some suggestions I think might be worth trying out.

[sayakpaul]

Could you please shed some light on the dataset you're using? Since it involves somewhat large-scale training, I won't be able to try things out myself. So, will try my best to provide some suggestions.

What happens when you try to overfit a single batch? Like do you see similar validation behavior, then? Also, on Step 1050, I see that the first two rows generate decent results but the final one doesn't. Have you tried evaluating that checkpoint in isolation on different images and seen anything weird?

[brycegoh]

Hi @sayakpaul, thanks for the reply. Actually I don't need to use the script in research project. The script in the examples/controlnet/train_controlnet_sdxl.py works as well.

However, I was reading the script and realise that the crops_coords_top_left_h and crops_coords_top_left_w are predefined as args to the script and default as 0. Doesn't that assume that all images in the dataset are of same resolution or are 1:1 aspect ratio?

The micro conditioning on image size is implemented here:

def compute_embeddings(batch, proportion_empty_prompts, text_encoders, tokenizers, is_train=True):
        original_size = (args.resolution, args.resolution)
        target_size = (args.resolution, args.resolution)
        crops_coords_top_left = (args.crops_coords_top_left_h, args.crops_coords_top_left_w) # predefined as args and defaults to 0
        prompt_batch = batch[args.caption_column]

        prompt_embeds, pooled_prompt_embeds = encode_prompt(
            prompt_batch, text_encoders, tokenizers, proportion_empty_prompts, is_train
        )
        add_text_embeds = pooled_prompt_embeds

        # Adapted from pipeline.StableDiffusionXLPipeline._get_add_time_ids
        add_time_ids = list(original_size + crops_coords_top_left + target_size)
        add_time_ids = torch.tensor([add_time_ids])

        prompt_embeds = prompt_embeds.to(accelerator.device)
        add_text_embeds = add_text_embeds.to(accelerator.device)
        add_time_ids = add_time_ids.repeat(len(prompt_batch), 1)
        add_time_ids = add_time_ids.to(accelerator.device, dtype=prompt_embeds.dtype)
        unet_added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}

        return {"prompt_embeds": prompt_embeds, **unet_added_cond_kwargs}

Therefore, am I right to assume that we should add the micro conditioning on image sizes as well? This means we need to track the crop coordinates and original size here before computing embeddings.

[brycegoh]

I attempted to add the image size micro condition.

For simplicity, the following snippet will only show relevant code that I added to add image size conditioning.

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[args.image_column]]

    conditioning_images = [image.convert("RGB") for image in examples[args.conditioning_image_column]]

    original_sizes = []
    all_images = []
    all_conditioning_images = []
    crop_top_lefts = []
    for image, conditioning_image in zip(images, conditioning_images):
        original_sizes.append((image.height, image.width))
        image = train_resize(image)
        conditioning_image = train_resize(conditioning_image)

        if args.random_flip and random.random() < 0.5:
            # flip
            image = train_flip(image)
        if args.center_crop:
            y1 = max(0, int(round((image.height - args.resolution) / 2.0)))
            x1 = max(0, int(round((image.width - args.resolution) / 2.0)))
            image = train_crop(image)
            conditioning_image = train_crop(conditioning_image)
        else:
            y1, x1, h, w = train_crop.get_params(image, (args.resolution, args.resolution))
            image = crop(image, y1, x1, h, w)
            conditioning_image = crop(conditioning_image, y1, x1, h, w)
        crop_top_left = (y1, x1)
        crop_top_lefts.append(crop_top_left)
        image = train_transforms(image)
        all_images.append(image)
        conditioning_image = conditioning_image_transforms(conditioning_image)
        all_conditioning_images.append(conditioning_image)
    
    examples["original_sizes"] = original_sizes
    examples["crop_top_lefts"] = crop_top_lefts
    examples["pixel_values"] = all_images
    examples["conditioning_pixel_values"] = all_conditioning_images
    return examples

with accelerator.main_process_first():
    dataset = dataset.with_transform(preprocess_train)
    
def encode_prompt(batch, text_encoders, tokenizers, proportion_empty_prompts, is_train=True):
    prompt_embeds_list = []
    prompt_batch = batch[args.caption_column]

    captions = []
    for caption in prompt_batch:
        if random.random() < proportion_empty_prompts:
            captions.append("")
        elif isinstance(caption, str):
            captions.append(caption)
        elif isinstance(caption, (list, np.ndarray)):
            # take a random caption if there are multiple
            captions.append(random.choice(caption) if is_train else caption[0])

    with torch.no_grad():
        for tokenizer, text_encoder in zip(tokenizers, text_encoders):
            text_inputs = tokenizer(
                captions,
                padding="max_length",
                max_length=tokenizer.model_max_length,
                truncation=True,
                return_tensors="pt",
            )
            text_input_ids = text_inputs.input_ids
            prompt_embeds = text_encoder(
                text_input_ids.to(text_encoder.device),
                output_hidden_states=True,
            )

            # We are only ALWAYS interested in the pooled output of the final text encoder
            pooled_prompt_embeds = prompt_embeds[0]
            prompt_embeds = prompt_embeds.hidden_states[-2]
            bs_embed, seq_len, _ = prompt_embeds.shape
            prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
            prompt_embeds_list.append(prompt_embeds)

    prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
    pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1)
    return {"prompt_embeds": prompt_embeds, "pooled_prompt_embeds": pooled_prompt_embeds}
    

compute_embeddings_fn = functools.partial(
    encode_prompt,
    text_encoders=text_encoders,
    tokenizers=tokenizers,
    proportion_empty_prompts=args.proportion_empty_prompts,
)
new_fingerprint = Hasher.hash(args) + "train"
train_dataset = train_dataset.map(compute_embeddings_fn, batched=True, new_fingerprint=new_fingerprint)

def collate_fn(examples):
      pixel_values = torch.stack([torch.tensor(example["pixel_values"]) for example in examples])
      pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
      
      conditioning_pixel_values = torch.stack([example["conditioning_pixel_values"] for example in examples])
      conditioning_pixel_values = conditioning_pixel_values.to(memory_format=torch.contiguous_format).to(dtype=weight_dtype)

      original_sizes = [example["original_sizes"] for example in examples]
      crop_top_lefts = [example["crop_top_lefts"] for example in examples]
      prompt_embeds = torch.stack([torch.tensor(example["prompt_embeds"]) for example in examples])
      pooled_prompt_embeds = torch.stack([torch.tensor(example["pooled_prompt_embeds"]) for example in examples])

      return {
          "pixel_values": pixel_values,
          "conditioning_pixel_values": conditioning_pixel_values,
          "prompt_embeds": prompt_embeds,
          "pooled_prompt_embeds": pooled_prompt_embeds,
          "original_sizes": original_sizes,
          "crop_top_lefts": crop_top_lefts,
      }

  train_dataloader = torch.utils.data.DataLoader(
      train_dataset,
      shuffle=True,
      collate_fn=collate_fn,
      batch_size=args.train_batch_size,
      num_workers=args.dataloader_num_workers,
  )
  
  for step, batch in enumerate(train_dataloader):
    def compute_time_ids(original_size, crops_coords_top_left):
        # Adapted from pipeline.StableDiffusionXLPipeline._get_add_time_ids
        target_size = (args.resolution, args.resolution)
        add_time_ids = list(original_size + crops_coords_top_left + target_size)
        add_time_ids = torch.tensor([add_time_ids])
        add_time_ids = add_time_ids.to(accelerator.device, dtype=weight_dtype)
        return add_time_ids
                
    add_time_ids = torch.cat(
        [compute_time_ids(s, c) for s, c in zip(batch["original_sizes"], batch["crop_top_lefts"])]
    )
    prompt_embeds = batch["prompt_embeds"].to(accelerator.device)
    unet_added_conditions = {"time_ids": add_time_ids.to(accelerator.device, dtype=prompt_embeds.dtype)}
    pooled_prompt_embeds = batch["pooled_prompt_embeds"].to(accelerator.device)
    unet_added_conditions.update({"text_embeds": pooled_prompt_embeds})
    
    down_block_res_samples, mid_block_res_sample = controlnet(
        noisy_latents,
        timesteps,
        encoder_hidden_states=prompt_embeds,
        added_cond_kwargs=unet_added_conditions,
        controlnet_cond=controlnet_image,
        return_dict=False,
    )

    # Predict the noise residual
    model_pred = unet(
        noisy_latents.to(dtype=weight_dtype),
        timesteps,
        encoder_hidden_states=prompt_embeds.to(dtype=weight_dtype),
        added_cond_kwargs=unet_added_conditions,
        down_block_additional_residuals=[
            sample.to(dtype=weight_dtype) for sample in down_block_res_samples
        ],
        mid_block_additional_residual=mid_block_res_sample.to(dtype=weight_dtype),
        return_dict=False,
    )[0]