diff --git a/lama_cleaner/app/src/App.tsx b/lama_cleaner/app/src/App.tsx index 1e50a7d..1c2e4c6 100644 --- a/lama_cleaner/app/src/App.tsx +++ b/lama_cleaner/app/src/App.tsx @@ -8,12 +8,14 @@ import { enableFileManagerState, fileState, isDisableModelSwitchState, + isEnableAutoSavingState, toastState, } from './store/Atoms' import { keepGUIAlive } from './utils' import Header from './components/Header/Header' import useHotKey from './hooks/useHotkey' import { + getEnableAutoSaving, getEnableFileManager, getIsDisableModelSwitch, isDesktop, @@ -34,6 +36,7 @@ function App() { const userInputImage = useInputImage() const setIsDisableModelSwitch = useSetRecoilState(isDisableModelSwitchState) const setEnableFileManager = useSetRecoilState(enableFileManagerState) + const setIsEnableAutoSavingState = useSetRecoilState(isEnableAutoSavingState) // Set Input Image useEffect(() => { @@ -66,7 +69,17 @@ function App() { setEnableFileManager(isEnabled === 'true') } fetchData2() - }, [setEnableFileManager, setIsDisableModelSwitch]) + + const fetchData3 = async () => { + const isEnabled = await getEnableAutoSaving().then(res => res.text()) + setIsEnableAutoSavingState(isEnabled === 'true') + } + fetchData3() + }, [ + setEnableFileManager, + setIsDisableModelSwitch, + setIsEnableAutoSavingState, + ]) // Dark Mode Hotkey useHotKey( diff --git a/lama_cleaner/app/src/adapters/inpainting.ts b/lama_cleaner/app/src/adapters/inpainting.ts index 5529a9c..ed94c30 100644 --- a/lama_cleaner/app/src/adapters/inpainting.ts +++ b/lama_cleaner/app/src/adapters/inpainting.ts @@ -122,6 +122,12 @@ export function getEnableFileManager() { }) } +export function getEnableAutoSaving() { + return fetch(`${API_ENDPOINT}/is_enable_auto_saving`, { + method: 'GET', + }) +} + export function switchModel(name: string) { const fd = new FormData() fd.append('name', name) diff --git a/lama_cleaner/app/src/components/Editor/Editor.tsx b/lama_cleaner/app/src/components/Editor/Editor.tsx index 724d3f1..38a0058 100644 --- a/lama_cleaner/app/src/components/Editor/Editor.tsx +++ b/lama_cleaner/app/src/components/Editor/Editor.tsx @@ -20,7 +20,6 @@ import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil' import { useWindowSize, useKey, useKeyPressEvent } from 'react-use' import inpaint, { downloadToOutput, - makeGif, postInteractiveSeg, } from '../../adapters/inpainting' import Button from '../shared/Button' @@ -45,12 +44,11 @@ import { imageWidthState, interactiveSegClicksState, isDiffusionModelsState, + isEnableAutoSavingState, isInpaintingState, isInteractiveSegRunningState, isInteractiveSegState, - isPaintByExampleState, isPix2PixState, - isSDState, negativePropmtState, propmtState, runManuallyState, @@ -184,6 +182,7 @@ export default function Editor() { const [redoCurLines, setRedoCurLines] = useState([]) const [redoLineGroups, setRedoLineGroups] = useState([]) const enableFileManager = useRecoilValue(enableFileManagerState) + const isEnableAutoSaving = useRecoilValue(isEnableAutoSavingState) const setImageWidth = useSetRecoilState(imageWidthState) const setImageHeight = useSetRecoilState(imageHeightState) @@ -1101,7 +1100,7 @@ export default function Editor() { if (file === undefined) { return } - if (enableFileManager && renders.length > 0) { + if ((enableFileManager || isEnableAutoSaving) && renders.length > 0) { try { downloadToOutput(renders[renders.length - 1], file.name, file.type) setToastState({ diff --git a/lama_cleaner/app/src/store/Atoms.tsx b/lama_cleaner/app/src/store/Atoms.tsx index 703a828..d0a52ec 100644 --- a/lama_cleaner/app/src/store/Atoms.tsx +++ b/lama_cleaner/app/src/store/Atoms.tsx @@ -41,6 +41,7 @@ interface AppState { disableShortCuts: boolean isInpainting: boolean isDisableModelSwitch: boolean + isEnableAutoSaving: boolean isInteractiveSeg: boolean isInteractiveSegRunning: boolean interactiveSegClicks: number[][] @@ -58,6 +59,7 @@ export const appState = atom({ disableShortCuts: false, isInpainting: false, isDisableModelSwitch: false, + isEnableAutoSaving: false, isInteractiveSeg: false, isInteractiveSegRunning: false, interactiveSegClicks: [], @@ -221,6 +223,18 @@ export const isDisableModelSwitchState = selector({ }, }) +export const isEnableAutoSavingState = selector({ + key: 'isEnableAutoSavingState', + get: ({ get }) => { + const app = get(appState) + return app.isEnableAutoSaving + }, + set: ({ get, set }, newValue: any) => { + const app = get(appState) + set(appState, { ...app, isEnableAutoSaving: newValue }) + }, +}) + export const croperState = atom({ key: 'croperState', default: { diff --git a/lama_cleaner/const.py b/lama_cleaner/const.py index 32370d9..401937a 100644 --- a/lama_cleaner/const.py +++ b/lama_cleaner/const.py @@ -59,7 +59,7 @@ Model download directory (by setting XDG_CACHE_HOME environment variable), by de """ OUTPUT_DIR_HELP = """ -Only required when --input is directory. Result images will be saved to output directory automatically. +Result images will be saved to output directory automatically without confirmation. """ INPUT_HELP = """ diff --git a/lama_cleaner/file_manager/file_manager.py b/lama_cleaner/file_manager/file_manager.py index 961016c..d23a6bc 100644 --- a/lama_cleaner/file_manager/file_manager.py +++ b/lama_cleaner/file_manager/file_manager.py @@ -14,7 +14,7 @@ from PIL import Image, ImageOps, PngImagePlugin from loguru import logger LARGE_ENOUGH_NUMBER = 100 -PngImagePlugin.MAX_TEXT_CHUNK = LARGE_ENOUGH_NUMBER * (1024 ** 2) +PngImagePlugin.MAX_TEXT_CHUNK = LARGE_ENOUGH_NUMBER * (1024**2) from .storage_backends import FilesystemStorageBackend from .utils import aspect_to_string, generate_filename, glob_img @@ -63,11 +63,11 @@ class FileManager(FileSystemEventHandler): if event.src_path == str(self.root_directory): logger.info(f"Image directory {event.src_path} modified") self.image_dir_filenames = self._media_names(self.root_directory) - self.modified_time['image'] = datetime.utcnow() + self.modified_time["image"] = datetime.utcnow() elif event.src_path == str(self.output_dir): logger.info(f"Output directory {event.src_path} modified") self.output_dir_filenames = self._media_names(self.output_dir) - self.modified_time['output'] = datetime.utcnow() + self.modified_time["output"] = datetime.utcnow() def init_app(self, app): if self.app is None: @@ -83,21 +83,15 @@ class FileManager(FileSystemEventHandler): app.extensions["thumbnail"] = self app.config.setdefault("THUMBNAIL_MEDIA_ROOT", self._default_root_directory) - app.config.setdefault("THUMBNAIL_MEDIA_THUMBNAIL_ROOT", self._default_thumbnail_directory) + app.config.setdefault( + "THUMBNAIL_MEDIA_THUMBNAIL_ROOT", self._default_thumbnail_directory + ) app.config.setdefault("THUMBNAIL_MEDIA_URL", self._default_root_url) - app.config.setdefault("THUMBNAIL_MEDIA_THUMBNAIL_URL", self._default_thumbnail_root_url) + app.config.setdefault( + "THUMBNAIL_MEDIA_THUMBNAIL_URL", self._default_thumbnail_root_url + ) app.config.setdefault("THUMBNAIL_DEFAULT_FORMAT", self._default_format) - def save_to_output_directory(self, image: np.ndarray, filename: str): - fp = Path(filename) - new_name = fp.stem + f"_{int(time.time())}" + fp.suffix - if image.shape[2] == 3: - image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) - elif image.shape[2] == 4: - image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA) - - cv2.imwrite(str(self.output_dir / new_name), image) - @property def root_directory(self): path = self.app.config["THUMBNAIL_MEDIA_ROOT"] @@ -137,14 +131,23 @@ class FileManager(FileSystemEventHandler): for name in names: path = os.path.join(directory, name) img = Image.open(path) - res.append({"name": name, "height": img.height, "width": img.width, "ctime": os.path.getctime(path)}) + res.append( + { + "name": name, + "height": img.height, + "width": img.width, + "ctime": os.path.getctime(path), + } + ) return res @property def thumbnail_url(self): return self.app.config["THUMBNAIL_MEDIA_THUMBNAIL_URL"] - def get_thumbnail(self, directory: Path, original_filename: str, width, height, **options): + def get_thumbnail( + self, directory: Path, original_filename: str, width, height, **options + ): storage = FilesystemStorageBackend(self.app) crop = options.get("crop", "fit") background = options.get("background") @@ -163,13 +166,19 @@ class FileManager(FileSystemEventHandler): thumbnail_size = (width, height) thumbnail_filename = generate_filename( - original_filename, aspect_to_string(thumbnail_size), crop, background, quality + original_filename, + aspect_to_string(thumbnail_size), + crop, + background, + quality, ) thumbnail_filepath = os.path.join( self.thumbnail_directory, original_path, thumbnail_filename ) - thumbnail_url = os.path.join(self.thumbnail_url, original_path, thumbnail_filename) + thumbnail_url = os.path.join( + self.thumbnail_url, original_path, thumbnail_filename + ) if storage.exists(thumbnail_filepath): return thumbnail_url, (width, height) @@ -183,7 +192,9 @@ class FileManager(FileSystemEventHandler): # get original image format options["format"] = options.get("format", image.format) - image = self._create_thumbnail(image, thumbnail_size, crop, background=background) + image = self._create_thumbnail( + image, thumbnail_size, crop, background=background + ) raw_data = self.get_raw_data(image, **options) storage.save(thumbnail_filepath, raw_data) diff --git a/lama_cleaner/parse_args.py b/lama_cleaner/parse_args.py index f9946b3..aaaddf3 100644 --- a/lama_cleaner/parse_args.py +++ b/lama_cleaner/parse_args.py @@ -170,15 +170,14 @@ def parse_args(): parser.error( f"invalid --input: {args.input} is a directory, --output-dir is required" ) - else: - output_dir = Path(args.output_dir) - if not output_dir.exists(): - logger.info(f"Creating output directory: {output_dir}") - output_dir.mkdir(parents=True) - else: - if not output_dir.is_dir(): - parser.error( - f"invalid --output-dir: {output_dir} is not a directory" - ) + + if args.output_dir is not None: + output_dir = Path(args.output_dir) + if not output_dir.exists(): + logger.info(f"Creating output directory: {output_dir}") + output_dir.mkdir(parents=True) + else: + if not output_dir.is_dir(): + parser.error(f"invalid --output-dir: {output_dir} is not a directory") return args diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py index 6ff0806..7fe31c2 100644 --- a/lama_cleaner/server.py +++ b/lama_cleaner/server.py @@ -83,11 +83,13 @@ CORS(app, expose_headers=["Content-Disposition"]) model: ModelManager = None thumb: FileManager = None +output_dir: str = None interactive_seg_model: InteractiveSeg = None device = None input_image_path: str = None is_disable_model_switch: bool = False is_enable_file_manager: bool = False +is_enable_auto_saving: bool = False is_desktop: bool = False @@ -124,11 +126,20 @@ def make_gif(): @app.route("/save_image", methods=["POST"]) def save_image(): - # all image in output directory + if output_dir is None: + return "--output-dir is None", 500 + input = request.files + filename = request.form["filename"] origin_image_bytes = input["image"].read() # RGB image, _ = load_img(origin_image_bytes) - thumb.save_to_output_directory(image, request.form["filename"]) + if image.shape[2] == 3: + image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + elif image.shape[2] == 4: + image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA) + + cv2.imwrite(os.path.join(output_dir, filename), image) + return "ok", 200 @@ -348,6 +359,12 @@ def get_is_enable_file_manager(): return res, 200 +@app.route("/is_enable_auto_saving") +def get_is_enable_auto_saving(): + res = "true" if is_enable_auto_saving else "false" + return res, 200 + + @app.route("/model_downloaded/") def model_downloaded(name): return str(model.is_downloaded(name)), 200 @@ -408,6 +425,12 @@ def main(args): global is_enable_file_manager global is_desktop global thumb + global output_dir + global is_enable_auto_saving + + output_dir = args.output_dir + if output_dir is not None: + is_enable_auto_saving = True device = torch.device(args.device) is_disable_model_switch = args.disable_model_switch