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
"vscode:/vscode.git/clone" did not exist on "01ec3db9324a6c2a3ac2bdc6b99ed729528a10e1"
Commit
6db777b3
authored
Mar 13, 2023
by
pythongosssss
Browse files
Added ability to save images to temp dir
parent
a50e1118
Changes
3
Show 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
sys
import
shutil
import
threading
import
asyncio
...
...
@@ -53,7 +54,14 @@ def hijack_progress(server):
return
v
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__"
:
cleanup_temp
()
loop
=
asyncio
.
new_event_loop
()
asyncio
.
set_event_loop
(
loop
)
server
=
server
.
PromptServer
(
loop
)
...
...
@@ -93,3 +101,4 @@ if __name__ == "__main__":
else
:
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:
class
SaveImage
:
def
__init__
(
self
):
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
def
INPUT_TYPES
(
s
):
return
{
"required"
:
{
"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"
},
}
...
...
@@ -791,7 +793,7 @@ class SaveImage:
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
):
prefix_len
=
len
(
filename_prefix
)
prefix
=
filename
[:
prefix_len
+
1
]
...
...
@@ -818,8 +820,21 @@ class SaveImage:
if
extra_pnginfo
is
not
None
:
for
x
in
extra_pnginfo
:
metadata
.
add_text
(
x
,
json
.
dumps
(
extra_pnginfo
[
x
]))
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
)
counter
+=
1
return
{
"ui"
:
{
"images"
:
paths
}
}
...
...
server.py
View file @
6db777b3
...
...
@@ -113,7 +113,7 @@ class PromptServer():
async
def
view_image
(
request
):
if
"file"
in
request
.
match_info
:
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
)
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