Ptv Dataset Loader

If you're poking at ComfyUI-Pt-Wrapper (the "no-code PyTorch training in ComfyUI" pack), this is the node you'll reach for first. It collapses two boring steps - download a classic torchvision dataset and wrap it in a DataLoader - into one box. Name a dataset, pick a batch size, hit run, and a PTDATALOADER pops out the other side ready to plug into a training node. That's it.

The pack, by the way, is a spin-off of the author's earlier ComfyUI-Data-Analysis, and it's genuinely obscure - no community noise around it, just a well-documented repo from HowToSD. Everything here is grounded in that repo, not in folklore.

How it works

Under the hood it's a thin wrapper around two PyTorch pieces: torchvision.datasets and torch.utils.data.DataLoader. The name field is a string that gets resolved via getattr(torchvision.datasets, name) - so you type MNIST, FashionMNIST, or CIFAR10 and you get the real class, not a lookup table of choices. It auto-applies a ToTensor() transform, downloads the data if you ask it to, and returns a ready-to-iterate DataLoader.

The inputs that matter

• name - the torchvision dataset class, e.g. MNIST, FashionMNIST, CIFAR10.
• download - set to True unless you already have it.
• root - where the dataset lands. Relative paths resolve under the pack's own datasets/ folder; an absolute path also works. Leave blank for the default.
• batch_size - samples per batch. This is the first knob you pull when you hit an out-of-memory error.
• shuffle - True shuffles each epoch, which you want for training.

The two optional text fields are where it gets fiddly. dataset_parameters takes a Python dict like {"train": True} for the train split or {"train": False} for the test set. load_parameters passes extra kwargs to the DataLoader.

Here's the trap: the shipped default for load_parameters is the string {"num_workers:1} - note the missing quote. That's malformed Python, so if you leave it alone the node dies with a SyntaxError the moment it hits ast.literal_eval. Fix it to {} or {"num_workers": 1} and you're fine.

Wiring it in

The single output is a PTDATALOADER, a custom type that only this pack's nodes understand. Feed it into the pack's training nodes (Pt Train Classification Model, Pt Evaluate Classification Model) or the bare Pt Data Loader node if you want to peek before training. It won't plug into stock ComfyUI nodes - that's expected.

Install

Install the whole pack once; every Ptv/Sp node comes with it.

cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper

Then restart ComfyUI. Or use ComfyUI Manager and search for "ComfyUI-Pt-Wrapper". The first pip install pulls the heavy ML stack from requirements.txt (torchvision, transformers, datasets, scikit-learn, and a pinned gensim==4.3.2) - give it a minute.

Common issues

• SyntaxError on first run - the load_parameters default, as above. Replace it.
• "Dataset 'X' is not available" - you typed a name torchvision doesn't have. The author only officially tests MNIST, FashionMNIST and CIFAR-10, so anything else is "might work, file an issue."
• OOM during training - lower batch_size.
• First run downloads the dataset (CIFAR-10 is ~170MB), so don't panic at the network spinner.

Start with FashionMNIST - it downloads fast, trains in seconds, and the pack ships ready-made train/eval workflows for it.