# ui/workspace.py # Contains functions for saving, loading, and managing workspace settings from files. import gradio as gr import json import os import traceback import tkinter as tk from tkinter import filedialog from PIL import Image # Import shared state and other managers from . import shared_state from . import metadata as metadata_manager # --- Constants --- # Define filenames and the default output folder at the module level. # These were moved from demo_gradio_svc.py outputs_folder = './outputs_svc/' SETTINGS_FILENAME = "goan_settings.json" UNLOAD_SAVE_FILENAME = "goan_unload_save.json" REFRESH_IMAGE_FILENAME = "goan_refresh_image.png" # --- Core Save/Load Logic --- def get_default_values_map(): """Returns a dictionary with the default values for all UI settings.""" return { 'prompt': '', 'n_prompt': '', 'total_second_length': 5.0, 'seed': -1, 'use_teacache': True, 'preview_frequency_ui': 5, 'segments_to_decode_csv': '', 'gs_ui': 10.0, 'gs_schedule_shape_ui': 'Off', 'gs_final_ui': 10.0, 'steps': 25, 'cfg': 1.0, 'latent_window_size': 9, 'gpu_memory_preservation': 6.0, 'use_fp32_transformer_output_ui': False, 'rs': 0.0, 'mp4_crf': 18, 'output_folder_ui': outputs_folder, } def save_settings_to_file(filepath, *ui_values_tuple): """Saves a tuple of UI values to a specified JSON file.""" settings_to_save = dict(zip(shared_state.ALL_TASK_UI_KEYS, ui_values_tuple)) try: with open(filepath, 'w', encoding='utf-8') as f: json.dump(settings_to_save, f, indent=4) gr.Info(f"Workspace saved to {filepath}") except Exception as e: gr.Warning(f"Error saving workspace: {e}") traceback.print_exc() def load_settings_from_file(filepath, return_updates=True): """Loads settings from a JSON file and returns Gradio updates or raw values.""" default_values = get_default_values_map() try: with open(filepath, 'r', encoding='utf-8') as f: loaded_settings = json.load(f) gr.Info(f"Loaded workspace from {filepath}") except Exception as e: gr.Warning(f"Could not load workspace from {filepath}: {e}") loaded_settings = {} final_settings = {**default_values, **loaded_settings} output_values = [final_settings.get(key, default_values.get(key)) for key in shared_state.ALL_TASK_UI_KEYS] # Type correction to prevent errors when loading from JSON for i, key in enumerate(shared_state.ALL_TASK_UI_KEYS): try: if key in ['seed', 'latent_window_size', 'steps', 'mp4_crf', 'preview_frequency_ui']: output_values[i] = int(output_values[i]) elif key in ['total_second_length', 'cfg', 'gs_ui', 'rs', 'gpu_memory_preservation', 'gs_final_ui']: output_values[i] = float(output_values[i]) elif key in ['use_teacache', 'use_fp32_transformer_output_ui']: output_values[i] = bool(output_values[i]) except (ValueError, TypeError): output_values[i] = default_values.get(key) return [gr.update(value=v) for v in output_values] if return_updates else output_values # --- NEW FUNCTION FOR INITIAL OUTPUT FOLDER LOADING --- def get_initial_output_folder_from_settings(): """ Attempts to load the 'output_folder_ui' value from UNLOAD_SAVE_FILENAME or SETTINGS_FILENAME. If not found or an error occurs, returns the module's default outputs_folder. Expands user home directory paths (e.g., '~') if present. """ # Use the module-level outputs_folder as the fallback default default_output_folder_path = outputs_folder # Prioritize the temporary unload save file, then the default settings file filename_to_check = None if os.path.exists(UNLOAD_SAVE_FILENAME): filename_to_check = UNLOAD_SAVE_FILENAME elif os.path.exists(SETTINGS_FILENAME): filename_to_check = SETTINGS_FILENAME if filename_to_check: try: with open(filename_to_check, 'r', encoding='utf-8') as f: settings = json.load(f) # If 'output_folder_ui' exists in the loaded settings, use it if 'output_folder_ui' in settings: # Expand user path (e.g., '~/.goan/outputs/') to its absolute form return os.path.expanduser(settings['output_folder_ui']) except Exception as e: # Log a warning if loading fails, but proceed with default print(f"Warning: Could not load 'output_folder_ui' from {filename_to_check} for initial path setup: {e}") traceback.print_exc() # Print full traceback for deeper debugging if needed # If no file found, or loading failed, or key not present, return the module's default return default_output_folder_path # --- UI Handler Functions --- def save_workspace(*ui_values_tuple): """Opens a file dialog to save the full workspace settings.""" root = tk.Tk() root.withdraw() file_path = filedialog.asksaveasfilename( defaultextension=".json", initialfile="goan_workspace.json", filetypes=[("JSON files", "*.json")] ) root.destroy() if file_path: save_settings_to_file(file_path, *ui_values_tuple) else: gr.Warning("Save cancelled.") def save_as_default_workspace(*ui_values_tuple): """Saves the current UI settings as the default startup configuration.""" gr.Info(f"Saving current settings as default to {SETTINGS_FILENAME}") save_settings_to_file(SETTINGS_FILENAME, *ui_values_tuple) def save_ui_and_image_for_refresh(*args_from_ui_controls_tuple): """Saves UI state and the current image to temporary files for session recovery.""" gallery_list = args_from_ui_controls_tuple[0] all_ui_values_tuple = args_from_ui_controls_tuple[1:] full_params_map = dict(zip(shared_state.ALL_TASK_UI_KEYS, all_ui_values_tuple)) settings_to_save = full_params_map.copy() if gallery_list and isinstance(gallery_list[0], (tuple, Image.Image)): pil_image = gallery_list[0][0] if isinstance(gallery_list[0], tuple) else gallery_list[0] try: creative_params = {k: full_params_map.get(k) for k in shared_state.CREATIVE_PARAM_KEYS} pil_image = metadata_manager.write_image_metadata(pil_image, creative_params) # Use the output folder path from the UI settings output_folder_path = full_params_map.get('output_folder_ui', outputs_folder) refresh_image_path = os.path.join(output_folder_path, REFRESH_IMAGE_FILENAME) pil_image.save(refresh_image_path) settings_to_save["refresh_image_path"] = refresh_image_path except Exception as e: gr.Warning(f"Could not save refresh image: {e}") # Save all settings to the unload file save_settings_to_file(UNLOAD_SAVE_FILENAME, *settings_to_save.values()) def load_workspace(): """Opens a file dialog to load a workspace from a JSON file.""" root = tk.Tk() root.withdraw() file_path = filedialog.askopenfilename(filetypes=[("JSON files", "*.json")]) root.destroy() return load_settings_from_file(file_path) if file_path else [gr.update()] * len(shared_state.ALL_TASK_UI_KEYS) def load_workspace_on_start(): """ Loads settings on app startup, prioritizing a temporary session file, then a default file, and finally falling back to hardcoded defaults. """ image_path_to_load = None settings_file = None if os.path.exists(UNLOAD_SAVE_FILENAME): settings_file = UNLOAD_SAVE_FILENAME try: with open(settings_file, 'r') as f: settings = json.load(f) # Check if the saved image path exists if "refresh_image_path" in settings and os.path.exists(settings["refresh_image_path"]): image_path_to_load = settings["refresh_image_path"] except Exception: pass # Ignore errors reading the temp file elif os.path.exists(SETTINGS_FILENAME): settings_file = SETTINGS_FILENAME if settings_file: print(f"Loading workspace from {settings_file}") ui_updates = load_settings_from_file(settings_file) if image_path_to_load: gr.Info("Restoring UI state and image from previous session.") if settings_file == UNLOAD_SAVE_FILENAME: os.remove(UNLOAD_SAVE_FILENAME) # Clean up temp file return [image_path_to_load] + ui_updates print("No workspace file found. Using default values.") default_vals = get_default_values_map() return [None] + [default_vals[key] for key in shared_state.ALL_TASK_UI_KEYS] def load_image_from_path(image_path): """Loads an image from a given path and deletes the temporary file.""" if image_path and os.path.exists(image_path): try: # The gallery component expects a list of (image, name) tuples. return gr.update(value=[(Image.open(image_path), "refresh_image")]) finally: os.remove(image_path) # Clean up the temp image return gr.update(value=None)