Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
chenpangpang
ComfyUI
Commits
6db777b3
Commit
6db777b3
authored
Mar 13, 2023
by
pythongosssss
Browse files
Added ability to save images to temp dir
parent
a50e1118
Changes
3
Hide whitespace changes
Inline
Side-by-side
Showing
3 changed files
with
28 additions
and
4 deletions
+28
-4
main.py
main.py
+9
-0
nodes.py
nodes.py
+18
-3
server.py
server.py
+1
-1
No files found.
main.py
View file @
6db777b3
import
os
import
os
import
sys
import
sys
import
shutil
import
threading
import
threading
import
asyncio
import
asyncio
...
@@ -53,7 +54,14 @@ def hijack_progress(server):
...
@@ -53,7 +54,14 @@ def hijack_progress(server):
return
v
return
v
setattr
(
tqdm
,
"update"
,
wrapped_func
)
setattr
(
tqdm
,
"update"
,
wrapped_func
)
def
cleanup_temp
():
temp_dir
=
os
.
path
.
join
(
os
.
path
.
dirname
(
os
.
path
.
realpath
(
__file__
)),
"temp"
)
if
os
.
path
.
exists
(
temp_dir
):
shutil
.
rmtree
(
temp_dir
)
if
__name__
==
"__main__"
:
if
__name__
==
"__main__"
:
cleanup_temp
()
loop
=
asyncio
.
new_event_loop
()
loop
=
asyncio
.
new_event_loop
()
asyncio
.
set_event_loop
(
loop
)
asyncio
.
set_event_loop
(
loop
)
server
=
server
.
PromptServer
(
loop
)
server
=
server
.
PromptServer
(
loop
)
...
@@ -93,3 +101,4 @@ if __name__ == "__main__":
...
@@ -93,3 +101,4 @@ if __name__ == "__main__":
else
:
else
:
loop
.
run_until_complete
(
run
(
server
,
address
=
address
,
port
=
port
,
verbose
=
not
dont_print
,
call_on_start
=
call_on_start
))
loop
.
run_until_complete
(
run
(
server
,
address
=
address
,
port
=
port
,
verbose
=
not
dont_print
,
call_on_start
=
call_on_start
))
cleanup_temp
()
nodes.py
View file @
6db777b3
...
@@ -775,12 +775,14 @@ class KSamplerAdvanced:
...
@@ -775,12 +775,14 @@ class KSamplerAdvanced:
class
SaveImage
:
class
SaveImage
:
def
__init__
(
self
):
def
__init__
(
self
):
self
.
output_dir
=
os
.
path
.
join
(
os
.
path
.
dirname
(
os
.
path
.
realpath
(
__file__
)),
"output"
)
self
.
output_dir
=
os
.
path
.
join
(
os
.
path
.
dirname
(
os
.
path
.
realpath
(
__file__
)),
"output"
)
self
.
temp_dir
=
os
.
path
.
join
(
os
.
path
.
dirname
(
os
.
path
.
realpath
(
__file__
)),
"temp"
)
@
classmethod
@
classmethod
def
INPUT_TYPES
(
s
):
def
INPUT_TYPES
(
s
):
return
{
"required"
:
return
{
"required"
:
{
"images"
:
(
"IMAGE"
,
),
{
"images"
:
(
"IMAGE"
,
),
"filename_prefix"
:
(
"STRING"
,
{
"default"
:
"ComfyUI"
})},
"filename_prefix"
:
(
"STRING"
,
{
"default"
:
"ComfyUI"
}),
"use_temp_file"
:
([
"yes"
,
"no"
],
{
"default"
:
"no"
}
),},
"hidden"
:
{
"prompt"
:
"PROMPT"
,
"extra_pnginfo"
:
"EXTRA_PNGINFO"
},
"hidden"
:
{
"prompt"
:
"PROMPT"
,
"extra_pnginfo"
:
"EXTRA_PNGINFO"
},
}
}
...
@@ -791,7 +793,7 @@ class SaveImage:
...
@@ -791,7 +793,7 @@ class SaveImage:
CATEGORY
=
"image"
CATEGORY
=
"image"
def
save_images
(
self
,
images
,
filename_prefix
=
"ComfyUI"
,
prompt
=
None
,
extra_pnginfo
=
None
):
def
save_images
(
self
,
images
,
filename_prefix
=
"ComfyUI"
,
use_temp_file
=
"no"
,
prompt
=
None
,
extra_pnginfo
=
None
):
def
map_filename
(
filename
):
def
map_filename
(
filename
):
prefix_len
=
len
(
filename_prefix
)
prefix_len
=
len
(
filename_prefix
)
prefix
=
filename
[:
prefix_len
+
1
]
prefix
=
filename
[:
prefix_len
+
1
]
...
@@ -818,8 +820,21 @@ class SaveImage:
...
@@ -818,8 +820,21 @@ class SaveImage:
if
extra_pnginfo
is
not
None
:
if
extra_pnginfo
is
not
None
:
for
x
in
extra_pnginfo
:
for
x
in
extra_pnginfo
:
metadata
.
add_text
(
x
,
json
.
dumps
(
extra_pnginfo
[
x
]))
metadata
.
add_text
(
x
,
json
.
dumps
(
extra_pnginfo
[
x
]))
file
=
f
"
{
filename_prefix
}
_
{
counter
:
05
}
_.png"
file
=
f
"
{
filename_prefix
}
_
{
counter
:
05
}
_.png"
img
.
save
(
os
.
path
.
join
(
self
.
output_dir
,
file
),
pnginfo
=
metadata
,
optimize
=
True
)
if
use_temp_file
==
"yes"
:
if
not
os
.
path
.
exists
(
self
.
temp_dir
):
os
.
makedirs
(
self
.
temp_dir
)
dir
=
self
.
temp_dir
else
:
dir
=
self
.
output_dir
img
.
save
(
os
.
path
.
join
(
dir
,
file
),
pnginfo
=
metadata
,
optimize
=
True
)
if
use_temp_file
==
"yes"
:
file
+=
"?type=temp"
paths
.
append
(
file
)
paths
.
append
(
file
)
counter
+=
1
counter
+=
1
return
{
"ui"
:
{
"images"
:
paths
}
}
return
{
"ui"
:
{
"images"
:
paths
}
}
...
...
server.py
View file @
6db777b3
...
@@ -113,7 +113,7 @@ class PromptServer():
...
@@ -113,7 +113,7 @@ class PromptServer():
async
def
view_image
(
request
):
async
def
view_image
(
request
):
if
"file"
in
request
.
match_info
:
if
"file"
in
request
.
match_info
:
type
=
request
.
rel_url
.
query
.
get
(
"type"
,
"output"
)
type
=
request
.
rel_url
.
query
.
get
(
"type"
,
"output"
)
if
type
!=
"output"
and
type
!=
"input"
:
if
type
not
in
[
"output"
,
"input"
,
"temp"
]
:
return
web
.
Response
(
status
=
400
)
return
web
.
Response
(
status
=
400
)
output_dir
=
os
.
path
.
join
(
os
.
path
.
dirname
(
os
.
path
.
realpath
(
__file__
)),
type
)
output_dir
=
os
.
path
.
join
(
os
.
path
.
dirname
(
os
.
path
.
realpath
(
__file__
)),
type
)
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment