Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
Overbill1683
Stable Diffusion Webui
提交
486bda9b
S
Stable Diffusion Webui
项目概览
Overbill1683
/
Stable Diffusion Webui
11 个月 前同步成功
通知
1771
Star
81
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
分析
仓库
DevOps
项目成员
Pages
S
Stable Diffusion Webui
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Pages
分析
分析
仓库分析
DevOps
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
提交
体验新版 GitCode,发现更多精彩内容 >>
未验证
提交
486bda9b
编写于
1月 13, 2023
作者:
A
AUTOMATIC1111
提交者:
GitHub
1月 13, 2023
浏览文件
操作
浏览文件
下载
差异文件
Merge pull request #6620 from guaneec/varsize_batch
Enable batch_size>1 for mixed-sized training
上级
d7aec59c
6be644fa
变更
1
隐藏空白更改
内联
并排
Showing
1 changed file
with
32 addition
and
4 deletion
+32
-4
modules/textual_inversion/dataset.py
modules/textual_inversion/dataset.py
+32
-4
未找到文件。
modules/textual_inversion/dataset.py
浏览文件 @
486bda9b
...
...
@@ -3,8 +3,10 @@ import numpy as np
import
PIL
import
torch
from
PIL
import
Image
from
torch.utils.data
import
Dataset
,
DataLoader
from
torch.utils.data
import
Dataset
,
DataLoader
,
Sampler
from
torchvision
import
transforms
from
collections
import
defaultdict
from
random
import
shuffle
,
choices
import
random
import
tqdm
...
...
@@ -45,12 +47,12 @@ class PersonalizedBase(Dataset):
assert
data_root
,
'dataset directory not specified'
assert
os
.
path
.
isdir
(
data_root
),
"Dataset directory doesn't exist"
assert
os
.
listdir
(
data_root
),
"Dataset directory is empty"
assert
batch_size
==
1
or
not
varsize
,
'variable img size must have batch size 1'
self
.
image_paths
=
[
os
.
path
.
join
(
data_root
,
file_path
)
for
file_path
in
os
.
listdir
(
data_root
)]
self
.
shuffle_tags
=
shuffle_tags
self
.
tag_drop_out
=
tag_drop_out
groups
=
defaultdict
(
list
)
print
(
"Preparing dataset..."
)
for
path
in
tqdm
.
tqdm
(
self
.
image_paths
):
...
...
@@ -103,13 +105,14 @@ class PersonalizedBase(Dataset):
if
include_cond
and
not
(
self
.
tag_drop_out
!=
0
or
self
.
shuffle_tags
):
with
devices
.
autocast
():
entry
.
cond
=
cond_model
([
entry
.
cond_text
]).
to
(
devices
.
cpu
).
squeeze
(
0
)
groups
[
image
.
size
].
append
(
len
(
self
.
dataset
))
self
.
dataset
.
append
(
entry
)
del
torchdata
del
latent_dist
del
latent_sample
self
.
length
=
len
(
self
.
dataset
)
self
.
groups
=
list
(
groups
.
values
())
assert
self
.
length
>
0
,
"No images have been found in the dataset."
self
.
batch_size
=
min
(
batch_size
,
self
.
length
)
self
.
gradient_step
=
min
(
gradient_step
,
self
.
length
//
self
.
batch_size
)
...
...
@@ -137,9 +140,34 @@ class PersonalizedBase(Dataset):
entry
.
latent_sample
=
shared
.
sd_model
.
get_first_stage_encoding
(
entry
.
latent_dist
).
to
(
devices
.
cpu
)
return
entry
class
GroupedBatchSampler
(
Sampler
):
def
__init__
(
self
,
data_source
:
PersonalizedBase
,
batch_size
:
int
):
n
=
len
(
data_source
)
self
.
groups
=
data_source
.
groups
self
.
len
=
n_batch
=
n
//
batch_size
expected
=
[
len
(
g
)
/
n
*
n_batch
*
batch_size
for
g
in
data_source
.
groups
]
self
.
base
=
[
int
(
e
)
//
batch_size
for
e
in
expected
]
self
.
n_rand_batches
=
nrb
=
n_batch
-
sum
(
self
.
base
)
self
.
probs
=
[
e
%
batch_size
/
nrb
/
batch_size
if
nrb
>
0
else
0
for
e
in
expected
]
self
.
batch_size
=
batch_size
def
__len__
(
self
):
return
self
.
len
def
__iter__
(
self
):
b
=
self
.
batch_size
for
g
in
self
.
groups
:
shuffle
(
g
)
batches
=
[]
for
g
in
self
.
groups
:
batches
.
extend
(
g
[
i
*
b
:(
i
+
1
)
*
b
]
for
i
in
range
(
len
(
g
)
//
b
))
for
_
in
range
(
self
.
n_rand_batches
):
rand_group
=
choices
(
self
.
groups
,
self
.
probs
)[
0
]
batches
.
append
(
choices
(
rand_group
,
k
=
b
))
shuffle
(
batches
)
yield
from
batches
class
PersonalizedDataLoader
(
DataLoader
):
def
__init__
(
self
,
dataset
,
latent_sampling_method
=
"once"
,
batch_size
=
1
,
pin_memory
=
False
):
super
(
PersonalizedDataLoader
,
self
).
__init__
(
dataset
,
shuffle
=
True
,
drop_last
=
True
,
batch_size
=
batch_size
,
pin_memory
=
pin_memory
)
super
(
PersonalizedDataLoader
,
self
).
__init__
(
dataset
,
batch_sampler
=
GroupedBatchSampler
(
dataset
,
batch_size
)
,
pin_memory
=
pin_memory
)
if
latent_sampling_method
==
"random"
:
self
.
collate_fn
=
collate_wrapper_random
else
:
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录