fix auto save to output_dir

This commit is contained in:
Qing 2024-02-09 15:17:28 +08:00
parent 314b2ae8c0
commit 131fa8300b
4 changed files with 11 additions and 9 deletions

View File

@ -167,6 +167,7 @@ class Api:
self.add_api_route("/api/v1/run_plugin_gen_image", self.api_run_plugin_gen_image, methods=["POST"])
self.add_api_route("/api/v1/samplers", self.api_samplers, methods=["GET"])
self.add_api_route("/api/v1/adjust_mask", self.api_adjust_mask, methods=["POST"])
self.add_api_route("/api/v1/save_image", self.api_save_image, methods=["POST"])
self.app.mount("/", StaticFiles(directory=WEB_APP_DIR, html=True), name="assets")
# fmt: on
@ -179,6 +180,12 @@ class Api:
def add_api_route(self, path: str, endpoint, **kwargs):
return self.app.add_api_route(path, endpoint, **kwargs)
def api_save_image(self, file: UploadFile):
filename = file.filename
origin_image_bytes = file.file.read()
with open(self.config.output_dir / filename, "wb") as fw:
fw.write(origin_image_bytes)
def api_current_model(self) -> ModelInfo:
return self.model_manager.current_model

View File

@ -144,7 +144,10 @@ def start(
device = check_device(device)
if input and not input.exists():
logger.error(f"invalid --input: {input} not exists")
exit()
exit(-1)
if input and input.is_dir() and not output_dir:
logger.error(f"invalid --output-dir: must be set when --input is a directory")
exit(-1)
if output_dir:
output_dir = output_dir.expanduser().absolute()
logger.info(f"Image will be saved to {output_dir}")

View File

@ -27,18 +27,11 @@ class FileManager:
self.thumbnail_directory.mkdir(parents=True)
# fmt: off
self.app.add_api_route("/api/v1/save_image", self.api_save_image, methods=["POST"])
self.app.add_api_route("/api/v1/medias", self.api_medias, methods=["GET"], response_model=List[MediasResponse])
self.app.add_api_route("/api/v1/media_file", self.api_media_file, methods=["GET"])
self.app.add_api_route("/api/v1/media_thumbnail_file", self.api_media_thumbnail_file, methods=["GET"])
# fmt: on
def api_save_image(self, file: UploadFile):
filename = file.filename
origin_image_bytes = file.file.read()
with open(self.output_dir / filename, "wb") as fw:
fw.write(origin_image_bytes)
def api_medias(self, tab: MediaTab) -> List[MediasResponse]:
img_dir = self._get_dir(tab)
return self._media_names(img_dir)

View File

@ -644,7 +644,6 @@ export function SettingsDialog() {
<Button
key={item}
variant="ghost"
disabled={item === TAB_PLUGINS && !removeBGEnabled}
onClick={() => setTab(item)}
className={cn(
tab === item ? "bg-muted " : "hover:bg-muted",