From f7d3f1e5caf675f1a0d1a172d382a0624b8d0165 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Tue, 20 Dec 2022 13:31:11 +0100
Subject: Dependency cleanup/upgrades

---
 data/csv.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

(limited to 'data/csv.py')

diff --git a/data/csv.py b/data/csv.py
index 6525e45..d400757 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -3,7 +3,6 @@ import torch
 import json
 import numpy as np
 from pathlib import Path
-import pytorch_lightning as pl
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader, random_split
 from torchvision import transforms
@@ -42,7 +41,7 @@ class CSVDataItem(NamedTuple):
     nprompt: str
 
 
-class CSVDataModule(pl.LightningDataModule):
+class CSVDataModule():
     def __init__(
             self,
             batch_size: int,
@@ -141,7 +140,7 @@ class CSVDataModule(pl.LightningDataModule):
             items = [
                 item
                 for item in items
-                if "mode" in item and self.mode in item["mode"]
+                if "mode" in item and self.mode in item["mode"].split(", ")
             ]
         items = self.prepare_items(template, expansions, items)
         items = self.filter_items(items)
-- 
cgit v1.2.3-54-g00ecf