Clear finetuning example

[hamza13-12]

I followed the tutorial for using a different scheduler and after loading all the other SD components, I was successfully able to use that pipeline for inference. However, I am having trouble finding a script that helps me to finetune the Unet module after having loaded the components like this. Particularly, I am struggling with writing the training loop so it would be extremely beneficial if the the tutorial for using a different scheduler could be carried on to show us how then to finetune a stable diffusion model

[tolgacangoz]

Hi @hamza13-12,
Could this page be what you are looking for?

[hamza13-12]

Hi @hamza13-12, Could this page be what you are looking for?

This is helpful but how do I how do I set up my own dataset into a format expected by this? I have two column in my dataset, image paths and image captions

[hamza13-12]

I have been trying to figure out how to do this.

https://colab.research.google.com/drive/1upKmc7BNkH7f0FPyRozjHDEYYtim8uh9?usp=sharing

You may view my progress on this colab link. But I have been struggling to identify how to finetune a diffusion model given a dataset of captions (prompts) and images. I'm thinking I could utilize ImageFolder on how to achieve this. But again, my question remains the same, as to how I could prepare my dataset for tinetuning suitable for a standard SD pipeline. Any help would be appreciated

[asomoza]

I don't understand your question completely, but currently your having problems with loading the images and the labels/captions if you have them in a local directory?

If is that your problem you need to create a custom dataset class that does this. Would a code example work for you? I have a rough one that I use but haven't really improved because right now training is not my main focus.

[hamza13-12]

A code example would be very helpful. I have the image paths along with the image captions in a dataframe. The paths are representative of the ones stored in my google drive. Now, I'm trying to load both these things into a dataset that would be suitable for SD and finetune it

[asomoza]

Alright, just a note, this is a quick dirty class I did to test my local images, use it as a guide, it works but I'm sure it does a lot of things wrong and haven't had the time learn more about training and test it better.

What you will need to focus is this:

        self.image_files = sorted([f for f in os.listdir(data_dir) if f.endswith((".jpg", ".jpeg", ".png"))])
        self.text_files = [os.path.splitext(f)[0] + ".txt" for f in self.image_files]

This class expects the images and the captions to be in the same directory and the captions file with the same name as the image. So if you have a dataframe with that information it shouldn't be hard to adapt it, just need make sure to open the correct image and caption so you can return the pixels values and the corresponding tokens.

I have an app to make the images so I know the exact sizes of the image variations, so I created those specific buckets for them.

import os

from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms


class LocalImageTextDataset(Dataset):
    def __init__(self, data_dir, tokenizer_one, tokenizer_two, size):
        self.data_dir = data_dir
        self.tokenizer_one = tokenizer_one
        self.tokenizer_two = tokenizer_two
        self.size = size
        self.transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize([0.5], [0.5]),
            ]
        )
        self.image_files = sorted([f for f in os.listdir(data_dir) if f.endswith((".jpg", ".jpeg", ".png"))])
        self.text_files = [os.path.splitext(f)[0] + ".txt" for f in self.image_files]

        self.buckets = {
            "1024x1024": [],
            "896x1152": [],
            "1152x896": [],
            "1344x768": [],
            "1536x704": [],
        }

        for i, image_file in enumerate(self.image_files):
            image = Image.open(os.path.join(self.data_dir, image_file))
            size = f"{image.width}x{image.height}"
            if size in self.buckets:
                self.buckets[size].append(i)

    def __getitem__(self, index):
        for _size, indices in self.buckets.items():
            if index < len(indices):
                image_file = self.image_files[indices[index]]
                break
            else:
                index -= len(indices)

        text_file = self.text_files[index]

        image = Image.open(os.path.join(self.data_dir, image_file))
        original_size = (image.height, image.width)

        train_resize = transforms.Resize(self.size, interpolation=transforms.InterpolationMode.BILINEAR)
        train_crop = transforms.CenterCrop(self.size)

        image = train_resize(image)

        y1 = max(0, int(round((image.height - self.size) / 2.0)))
        x1 = max(0, int(round((image.width - self.size) / 2.0)))
        image = train_crop(image)
        crop_top_left = (y1, x1)
        image = self.transform(image)

        with open(os.path.join(self.data_dir, text_file), "r", encoding="utf-8") as f:
            text = f.read()

        tokens_one = self.tokenize_captions(text, self.tokenizer_one).squeeze(dim=0)
        tokens_two = self.tokenize_captions(text, self.tokenizer_two).squeeze(dim=0)

        return {
            "pixel_values": image,
            "input_ids_one": tokens_one,
            "input_ids_two": tokens_two,
            "original_size": original_size,
            "crop_top_left": crop_top_left,
        }

    def tokenize_captions(self, caption, tokenizer):
        text_inputs = tokenizer(
            caption,
            padding="max_length",
            max_length=tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )
        return text_inputs.input_ids

    def __len__(self):
        return len(self.image_files)

[sayakpaul]

Thanks for this! Below is just a suggestion and NOTHING else.

About

        for i, image_file in enumerate(self.image_files):
            image = Image.open(os.path.join(self.data_dir, image_file))
            size = f"{image.width}x{image.height}"
            if size in self.buckets:
                self.buckets[size].append(i)

We could also keep track of this before initializing the dataloader as it could clog it depending on the number of samples you have. So, we could maybe maintain a mapping (between the image filename and its original resolution) in a thread-safe parallel manner. We could then use that mapping here to initialize the dataloader.

But as mentioned this is just an improvement suggestion. Your implementation is already pristine as a starting reference.

[asomoza]

thanks for that suggestion, I'll incorporate it in my training code. Also with this I noticed that even though I used the buckets, the images get resized and cropped to the same size, so I'll need to fix that too. Going to do some quick tests today and I'll update the code.

[sayakpaul]

This should be an issue. This should be transferred to the Discussions.