Merge pull request #736 from JeLuF/enfdir

Enforce an autosave directory

ui/easydiffusion/app.py

@@ -4,6 +4,7 @@ import sys
 import json
 import traceback
 import logging
+import shlex
 from rich.logging import RichHandler
 
 from sdkit.utils import log as sdkit_log # hack, so we can overwrite the log config
@@ -36,6 +37,7 @@ CORE_UI_PLUGINS_DIR = os.path.abspath(os.path.join(SD_UI_DIR, 'plugins', 'ui'))
 UI_PLUGINS_SOURCES = ((CORE_UI_PLUGINS_DIR, 'core'), (USER_UI_PLUGINS_DIR, 'user'))
 
 OUTPUT_DIRNAME = "Stable Diffusion UI" # in the user's home folder
+PRESERVE_CONFIG_VARS = ['FORCE_FULL_PRECISION']
 TASK_TTL = 15 * 60 # Discard last session's task timeout
 APP_CONFIG_DEFAULTS = {
     # auto: selects the cuda device with the most free memory, cuda: use the currently active cuda device.
@@ -88,6 +90,11 @@ def setConfig(config):
         bind_ip = '0.0.0.0' if config['net']['listen_to_network'] else '127.0.0.1'
         config_bat.append(f"@set SD_UI_BIND_IP={bind_ip}")
 
+        # Preserve these variables if they are set
+        for var in PRESERVE_CONFIG_VARS:
+            if os.getenv(var) is not None:
+                config_bat.append(f"@set {var}={os.getenv(var)}")
+
         if len(config_bat) > 0:
             with open(config_bat_path, 'w', encoding='utf-8') as f:
                 f.write('\r\n'.join(config_bat))
@@ -105,6 +112,11 @@ def setConfig(config):
         bind_ip = '0.0.0.0' if config['net']['listen_to_network'] else '127.0.0.1'
         config_sh.append(f"export SD_UI_BIND_IP={bind_ip}")
 
+        # Preserve these variables if they are set
+        for var in PRESERVE_CONFIG_VARS:
+            if os.getenv(var) is not None:
+                config_bat.append(f'export {var}="{shlex.quote(os.getenv(var))}"')
+
         if len(config_sh) > 1:
             with open(config_sh_path, 'w', encoding='utf-8') as f:
                 f.write('\n'.join(config_sh))

ui/easydiffusion/server.py

@@ -132,10 +132,14 @@ def read_web_data_internal(key:str=None):
         return JSONResponse(app.getConfig(), headers=NOCACHE_HEADERS)
     elif key == 'system_info':
         config = app.getConfig()
+
+        output_dir = config.get('force_save_path', os.path.join(os.path.expanduser("~"), app.OUTPUT_DIRNAME))
+
         system_info = {
             'devices': task_manager.get_devices(),
             'hosts': app.getIPConfig(),
-            'default_output_dir': os.path.join(os.path.expanduser("~"), app.OUTPUT_DIRNAME),
+            'default_output_dir': output_dir,
+            'enforce_output_dir': ('force_save_path' in config),
         }
         system_info['devices']['config'] = config.get('render_devices', "auto")
         return JSONResponse(system_info, headers=NOCACHE_HEADERS)
@@ -165,6 +169,11 @@ def render_internal(req: dict):
         render_req: GenerateImageRequest = GenerateImageRequest.parse_obj(req)
         task_data: TaskData = TaskData.parse_obj(req)
 
+        # Overwrite user specified save path
+        config=app.getConfig
+        if 'force_save_path' in config:
+            task_data.save_to_disk_path = config['force_save_path']
+
         render_req.init_image_mask = req.get('mask') # hack: will rename this in the HTTP API in a future revision
 
         app.save_to_config(task_data.use_stable_diffusion_model, task_data.use_vae_model, task_data.use_hypernetwork_model, task_data.vram_usage_level)

ui/media/js/parameters.js

@@ -329,9 +329,9 @@ autoPickGPUsField.addEventListener('click', function() {
     gpuSettingEntry.style.display = (this.checked ? 'none' : '')
 })
 
-async function setDiskPath(defaultDiskPath) {
+async function setDiskPath(defaultDiskPath, force=false) {
     var diskPath = getSetting("diskPath")
-    if (diskPath == '' || diskPath == undefined || diskPath == "undefined") {
+    if (force || diskPath == '' || diskPath == undefined || diskPath == "undefined") {
         setSetting("diskPath", defaultDiskPath)
     }
 }
@@ -407,7 +407,13 @@ async function getSystemInfo() {
 
         setDeviceInfo(devices)
         setHostInfo(res['hosts'])
-        setDiskPath(res['default_output_dir'])
+        let force = false
+        if (res['enforce_output_dir'] !== undefined) {
+            force = res['enforce_output_dir']
+            saveToDiskField.disabled = force
+            diskPathField.disabled = force
+        }
+        setDiskPath(res['default_output_dir'], force)
     } catch (e) {
         console.log('error fetching devices', e)
     }