theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-24 10:18:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

support json dataset

Browse Source
This commit is contained in:
Artiprocher
2025-07-02 20:07:16 +08:00
parent 9bb51fe879
commit 1363a0559f
1 changed files with 13 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
diffsynth/trainers/utils.py
View File

@@ -1,4 +1,4 @@
import imageio, os, torch, warnings, torchvision, argparse import imageio, os, torch, warnings, torchvision, argparse, json
from peft import LoraConfig, inject_adapter_in_model from peft import LoraConfig, inject_adapter_in_model
from PIL import Image from PIL import Image
import pandas as pd import pandas as pd
@@ -48,9 +48,14 @@ class ImageDataset(torch.utils.data.Dataset):
print("No metadata. Trying to generate it.") print("No metadata. Trying to generate it.")
metadata = self.generate_metadata(base_path) metadata = self.generate_metadata(base_path)
print(f"{len(metadata)} lines in metadata.") print(f"{len(metadata)} lines in metadata.")
self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
elif metadata_path.endswith(".json"):
with open(metadata_path, "r") as f:
metadata = json.load(f)
self.data = metadata
else: else:
metadata = pd.read_csv(metadata_path) metadata = pd.read_csv(metadata_path)
self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))] self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
def generate_metadata(self, folder): def generate_metadata(self, folder):
@@ -177,9 +182,14 @@ class VideoDataset(torch.utils.data.Dataset):
print("No metadata. Trying to generate it.") print("No metadata. Trying to generate it.")
metadata = self.generate_metadata(base_path) metadata = self.generate_metadata(base_path)
print(f"{len(metadata)} lines in metadata.") print(f"{len(metadata)} lines in metadata.")
self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
elif metadata_path.endswith(".json"):
with open(metadata_path, "r") as f:
metadata = json.load(f)
self.data = metadata
else: else:
metadata = pd.read_csv(metadata_path) metadata = pd.read_csv(metadata_path)
self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))] self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
def generate_metadata(self, folder): def generate_metadata(self, folder):