diffengine.datasets.hf_controlnet_datasets

Module Contents

Classes

HFControlNetDataset

Dataset for huggingface datasets.
class diffengine.datasets.hf_controlnet_datasets.HFControlNetDataset(dataset, image_column='image', condition_column='condition', caption_column='text', csv='metadata.csv', pipeline=(), cache_dir=None)[source]

Bases: torch.utils.data.Dataset

Dataset for huggingface datasets.

Args:

dataset (str): Dataset name or path to dataset. image_column (str): Image column name. Defaults to ‘image’. condition_column (str): Condition column name for ControlNet.

Defaults to ‘condition’.

caption_column (str): Caption column name. Defaults to ‘text’. csv (str): Caption csv file name when loading local folder.

Defaults to ‘metadata.csv’.

pipeline (Sequence): Processing pipeline. Defaults to an empty tuple. cache_dir (str, optional): The directory where the downloaded datasets

will be stored.Defaults to None.

__len__()[source]

Get the length of dataset.

Returns:

int

Return type:

The length of filtered dataset.

__getitem__(idx)[source]

Get item.

Get the idx-th image and data information of dataset after ``self.pipeline`.

Args:

idx (int): The index of self.data_list.

Returns:

dict: The idx-th image and data information of dataset after self.pipeline.

Parameters:

idx (int) –

Return type:

dict

Parameters:

  • dataset (str) –

  • image_column (str) –

  • condition_column (str) –

  • caption_column (str) –

  • csv (str) –

  • pipeline (collections.abc.Sequence) –

  • cache_dir (str | None) –