It is very difficult to train ControlNet using the original dataset.

[mattyamonaca]

The current Controlnet training code seems to be designed specifically for use with ControlNet-Hed.
In particular, there are few instructions on how to prepare the data set, and you need to read the code and modify it for your own use in order to prepare the data.

Therefore, I created a simple script to convert HF's ControlNet training dataset (e.g., the training format used in Diffusers) for use with PixArt-alpha.

from datasets import load_dataset
import os
from diffusers.models import AutoencoderKL
import numpy as np
import torchvision.transforms as transforms
import torch
from tqdm import tqdm
import json
dataset = load_dataset('ishapuri/human_animal_poses')
data_root = "./data/Openpose-Controlnet"

os.makedirs(data_root, exist_ok=True)
os.makedirs(os.path.join(data_root, "image"), exist_ok=True)
os.makedirs(os.path.join(data_root, "conditioning"), exist_ok=True)


data_info = []
transform = transforms.ToTensor()
vae = AutoencoderKL.from_pretrained("output/pretrained_models/sd-vae-ft-ema")
vae.eval()
vae.requires_grad_(False)
vae.to("cuda")


for idx, data in tqdm(enumerate(dataset["train"])):
    caption = data["caption"]
    image = data["original_image"]
    conditioning = data["condtioning_image"]
    ratio = image.height/image.width

    info = {
        "height": image.height,
        "width": image.width,
        "ratio": ratio,
        "path": f"image_{idx}.png",
        "prompt": caption 
    }

    image = image.resize((512,512))
    image.save(os.path.join(data_root, "image", f"image_{idx}.png"))
    
    conditioning = conditioning.resize((512, 512))
    conditioning = transform(conditioning).unsqueeze(0).cuda()

    posterior = vae.encode(conditioning).latent_dist
    conditioning = torch.cat([posterior.mean, posterior.std], dim=1).cpu().numpy()
    for p, o in zip([f"image_{idx}.png"], conditioning ):
        save = os.path.join(data_root, 'conditioning/') + p.replace('.png', '.npz')

        os.makedirs(os.path.dirname(save), exist_ok=True)
        np.savez_compressed(save, o)
    data_info.append(info)

with open(os.path.join(data_root, f"data_info.json"), 'w') as f:
    json.dump(data_info, f, indent=4)

If necessary, I will refactor the code for more general use and submit a PR to be added to the "tools" directory.
Additionally, I can suggest a more generic "datasets/pixart-control.py".

[darius522]

Thanks so much for providing this @mattyamonaca . Currently this line:
vae = AutoencoderKL.from_pretrained("output/pretrained_models/sd-vae-ft-ema")

Fails for me as HF's from_pretrained expects a safetensors format, which I don't currently have access to. When running download.py all I get is a .bin file for the sd-vae-ft-ema model. Should I look somewhere else for this file ?

[darius522]

As an FYI, it looks like the newer version of diffusers stopped supporting the .bin format altogether. Downgrading to 0.28.1 actually "fixed" the issue for me. Something to keep in mind ...

[mattyamonaca]

OK.
I'll prepare a working solution and submit a PR.
Thank you for sharing the information.

[raulc0399]

@mattyamonaca i have written a convert script to convert the HED model to safetensors and implemented a control net pipeline for HF diffuser that uses that. here the PR for that:
#158

i also implemented the training script for HF pixart controlnet. here is the PR for that:
#164

both for pixart alpha.

you can try the converted dataset with that, this will train and save only the controlnet layers and you can use the provided script to generate validation images.

can you share the converted dataset?

thanks.