From 07bdefd5a25381da9cb8ed649f7b910ad990f7c5 Mon Sep 17 00:00:00 2001 From: iconben <9401905+iconben@users.noreply.github.com> Date: Sun, 5 Apr 2026 05:11:14 +0000 Subject: [PATCH 1/3] feat: enforce max constraints on image generation limits Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> --- patch_cli_info.py | 58 ++++++++++++++++++++++++++++++++++++ patch_cli_info_fix.py | 11 +++++++ patch_cli_info_fix2.py | 32 ++++++++++++++++++++ patch_cli_info_fix3.py | 26 ++++++++++++++++ patch_cli_info_fix4.py | 34 +++++++++++++++++++++ patch_cli_info_imports.py | 27 +++++++++++++++++ patch_cli_info_imports_2.py | 11 +++++++ patch_cli_limits.py | 34 +++++++++++++++++++++ patch_js_limits.py | 50 +++++++++++++++++++++++++++++++ patch_js_limits_fix.py | 52 ++++++++++++++++++++++++++++++++ patch_mcp_limits.py | 58 ++++++++++++++++++++++++++++++++++++ patch_mcp_limits_schema.py | 26 ++++++++++++++++ patch_paths.py | 27 +++++++++++++++++ patch_server_info.py | 27 +++++++++++++++++ patch_server_info_fix.py | 30 +++++++++++++++++++ patch_server_info_fix2.py | 27 +++++++++++++++++ patch_server_limits.py | 31 +++++++++++++++++++ patch_test_cli_info.py | 32 ++++++++++++++++++++ src/zimage/cli.py | 30 +++++++++++++++++++ src/zimage/mcp_server.py | 22 ++++++++++++-- src/zimage/paths.py | 3 ++ src/zimage/server.py | 27 +++++++++++++++++ src/zimage/static/js/main.js | 29 ++++++++++++++++++ test_paths.py | 16 ++++++++++ test_server_info.py | 12 ++++++++ tests/test_cli_info.py | 10 +++++++ 26 files changed, 739 insertions(+), 3 deletions(-) create mode 100644 patch_cli_info.py create mode 100644 patch_cli_info_fix.py create mode 100644 patch_cli_info_fix2.py create mode 100644 patch_cli_info_fix3.py create mode 100644 patch_cli_info_fix4.py create mode 100644 patch_cli_info_imports.py create mode 100644 patch_cli_info_imports_2.py create mode 100644 patch_cli_limits.py create mode 100644 patch_js_limits.py create mode 100644 patch_js_limits_fix.py create mode 100644 patch_mcp_limits.py create mode 100644 patch_mcp_limits_schema.py create mode 100644 patch_paths.py create mode 100644 patch_server_info.py create mode 100644 patch_server_info_fix.py create mode 100644 patch_server_info_fix2.py create mode 100644 patch_server_limits.py create mode 100644 patch_test_cli_info.py create mode 100644 test_paths.py create mode 100644 test_server_info.py diff --git a/patch_cli_info.py b/patch_cli_info.py new file mode 100644 index 0000000..e4c3897 --- /dev/null +++ b/patch_cli_info.py @@ -0,0 +1,58 @@ +import re + +with open('src/zimage/cli.py', 'r') as f: + content = f.read() + +# 1. Update collect_info +old_collect = """ "paths": { + "module_file": str(Path(__file__).resolve()), + "config_path": str(get_config_path().resolve()), + "data_dir": str(get_data_dir().resolve()), + "outputs_dir": str(get_outputs_dir().resolve()), + "loras_dir": str(get_loras_dir().resolve()), + "db_path": str(get_db_path().resolve()), + }, + "hardware": hardware, + }""" + +new_collect = """ "paths": { + "module_file": str(Path(__file__).resolve()), + "config_path": str(get_config_path().resolve()), + "data_dir": str(get_data_dir().resolve()), + "outputs_dir": str(get_outputs_dir().resolve()), + "loras_dir": str(get_loras_dir().resolve()), + "db_path": str(get_db_path().resolve()), + }, + "constraints": { + "max_steps": load_config().get("max_steps", 50), + "max_width": load_config().get("max_width", 4096), + "max_height": load_config().get("max_height", 4096), + }, + "hardware": hardware, + }""" +content = content.replace(old_collect, new_collect) + + +# 2. Update format_info_text +old_format = """ f" Outputs Dir: {info['paths']['outputs_dir']}", + f" LoRAs Dir: {info['paths']['loras_dir']}", + f" DB Path: {info['paths']['db_path']}", + "", + "Hardware:" + ]""" + +new_format = """ f" Outputs Dir: {info['paths']['outputs_dir']}", + f" LoRAs Dir: {info['paths']['loras_dir']}", + f" DB Path: {info['paths']['db_path']}", + "", + "Constraints:", + f" Max Steps: {info['constraints']['max_steps']}", + f" Max Width: {info['constraints']['max_width']}", + f" Max Height: {info['constraints']['max_height']}", + "", + "Hardware:" + ]""" +content = content.replace(old_format, new_format) + +with open('src/zimage/cli.py', 'w') as f: + f.write(content) diff --git a/patch_cli_info_fix.py b/patch_cli_info_fix.py new file mode 100644 index 0000000..8551c38 --- /dev/null +++ b/patch_cli_info_fix.py @@ -0,0 +1,11 @@ +import re + +with open('src/zimage/cli.py', 'r') as f: + content = f.read() + +# Make sure constraints are actually output (I noticed they didn't show up in the output) +# Let's check where they should be inserted. +# Ah, I replaced "Environment Overrides:" with "Hardware:", I might have lost "Environment Overrides:". + +with open('src/zimage/cli.py', 'w') as f: + f.write(content) diff --git a/patch_cli_info_fix2.py b/patch_cli_info_fix2.py new file mode 100644 index 0000000..a1bdd7c --- /dev/null +++ b/patch_cli_info_fix2.py @@ -0,0 +1,32 @@ +import re + +with open('src/zimage/cli.py', 'r') as f: + content = f.read() + +old_str = """ f" Outputs Dir: {info['paths']['outputs_dir']}", + f" LoRAs Dir: {info['paths']['loras_dir']}", + f" DB Path: {info['paths']['db_path']}", + "", + "Environment Overrides:" + ]""" + +new_str = """ f" Outputs Dir: {info['paths']['outputs_dir']}", + f" LoRAs Dir: {info['paths']['loras_dir']}", + f" DB Path: {info['paths']['db_path']}", + "", + "Constraints:", + f" Max Steps: {info['constraints']['max_steps']}", + f" Max Width: {info['constraints']['max_width']}", + f" Max Height: {info['constraints']['max_height']}", + "", + "Environment Overrides:" + ]""" + +if old_str in content: + content = content.replace(old_str, new_str) +else: + # My previous replacement might have matched something else or failed + pass + +with open('src/zimage/cli.py', 'w') as f: + f.write(content) diff --git a/patch_cli_info_fix3.py b/patch_cli_info_fix3.py new file mode 100644 index 0000000..8cd8e4a --- /dev/null +++ b/patch_cli_info_fix3.py @@ -0,0 +1,26 @@ +import re + +with open('src/zimage/cli.py', 'r') as f: + content = f.read() + +old_str = """ f" Outputs Dir: {info['paths']['outputs_dir']}", + f" LoRAs Dir: {info['paths']['loras_dir']}", + f" DB Path: {info['paths']['db_path']}", + "", + "Environment Overrides:""" + +new_str = """ f" Outputs Dir: {info['paths']['outputs_dir']}", + f" LoRAs Dir: {info['paths']['loras_dir']}", + f" DB Path: {info['paths']['db_path']}", + "", + "Constraints:", + f" Max Steps: {info['constraints']['max_steps']}", + f" Max Width: {info['constraints']['max_width']}", + f" Max Height: {info['constraints']['max_height']}", + "", + "Environment Overrides:""" + +content = content.replace(old_str, new_str) + +with open('src/zimage/cli.py', 'w') as f: + f.write(content) diff --git a/patch_cli_info_fix4.py b/patch_cli_info_fix4.py new file mode 100644 index 0000000..0f88dee --- /dev/null +++ b/patch_cli_info_fix4.py @@ -0,0 +1,34 @@ +with open('src/zimage/cli.py', 'r') as f: + content = f.read() + +# Let's see what collect_info returns currently +old_collect = """ "paths": { + "module_file": str(Path(__file__).resolve()), + "config_path": str(get_config_path().resolve()), + "data_dir": str(get_data_dir().resolve()), + "outputs_dir": str(get_outputs_dir().resolve()), + "loras_dir": str(get_loras_dir().resolve()), + "db_path": str(get_db_path().resolve()), + }, + "env_overrides": {""" + +new_collect = """ "paths": { + "module_file": str(Path(__file__).resolve()), + "config_path": str(get_config_path().resolve()), + "data_dir": str(get_data_dir().resolve()), + "outputs_dir": str(get_outputs_dir().resolve()), + "loras_dir": str(get_loras_dir().resolve()), + "db_path": str(get_db_path().resolve()), + }, + "constraints": { + "max_steps": load_config().get("max_steps", 50), + "max_width": load_config().get("max_width", 4096), + "max_height": load_config().get("max_height", 4096), + }, + "env_overrides": {""" + +if old_collect in content: + content = content.replace(old_collect, new_collect) + +with open('src/zimage/cli.py', 'w') as f: + f.write(content) diff --git a/patch_cli_info_imports.py b/patch_cli_info_imports.py new file mode 100644 index 0000000..fa61d93 --- /dev/null +++ b/patch_cli_info_imports.py @@ -0,0 +1,27 @@ +import re + +with open('src/zimage/cli.py', 'r') as f: + content = f.read() + +# I need to add load_config to the imports from paths +content = content.replace( + """get_outputs_dir, + get_loras_dir, + get_db_path,""", + """get_outputs_dir, + get_loras_dir, + get_db_path, + load_config,""" +) +content = content.replace( + """get_outputs_dir, + get_loras_dir, + get_db_path,""", + """get_outputs_dir, + get_loras_dir, + get_db_path, + load_config,""" +) + +with open('src/zimage/cli.py', 'w') as f: + f.write(content) diff --git a/patch_cli_info_imports_2.py b/patch_cli_info_imports_2.py new file mode 100644 index 0000000..de42608 --- /dev/null +++ b/patch_cli_info_imports_2.py @@ -0,0 +1,11 @@ +import re + +with open('src/zimage/cli.py', 'r') as f: + content = f.read() + +content = content.replace("get_db_path,\n get_config_path,\n", "get_db_path,\n get_config_path,\n load_config,\n") +content = content.replace("get_db_path,\n get_config_path,\n", "get_db_path,\n get_config_path,\n load_config,\n") +content = content.replace("get_db_path,\n get_config_path,\n", "get_db_path,\n get_config_path,\n load_config,\n") + +with open('src/zimage/cli.py', 'w') as f: + f.write(content) diff --git a/patch_cli_limits.py b/patch_cli_limits.py new file mode 100644 index 0000000..c5ca3e7 --- /dev/null +++ b/patch_cli_limits.py @@ -0,0 +1,34 @@ +with open('src/zimage/cli.py', 'r') as f: + content = f.read() + +old_run_gen = """def run_generation(args): + generate_image, save_image, record_generation = _load_generation_modules() + logger.info(f"DEBUG: cwd: {Path.cwd().resolve()}") + + # Ensure width/height are multiples of 16""" + +new_run_gen = """def run_generation(args): + generate_image, save_image, record_generation = _load_generation_modules() + logger.info(f"DEBUG: cwd: {Path.cwd().resolve()}") + + # Check constraints + config = load_config() + max_steps = config.get("max_steps", 50) + max_width = config.get("max_width", 4096) + max_height = config.get("max_height", 4096) + + if args.steps > max_steps: + log_error(f"Requested steps ({args.steps}) exceeds the maximum allowed ({max_steps}).") + sys.exit(1) + if args.width > max_width: + log_error(f"Requested width ({args.width}) exceeds the maximum allowed ({max_width}).") + sys.exit(1) + if args.height > max_height: + log_error(f"Requested height ({args.height}) exceeds the maximum allowed ({max_height}).") + + # Ensure width/height are multiples of 16""" + +content = content.replace(old_run_gen, new_run_gen) + +with open('src/zimage/cli.py', 'w') as f: + f.write(content) diff --git a/patch_js_limits.py b/patch_js_limits.py new file mode 100644 index 0000000..5589832 --- /dev/null +++ b/patch_js_limits.py @@ -0,0 +1,50 @@ +with open('src/zimage/static/js/main.js', 'r') as f: + content = f.read() + +# Let's insert a call to /info at startup and update DOM inputs +old_init = """ // Initialize components + document.addEventListener("DOMContentLoaded", () => { + // Initialize tooltips + const tooltipTriggerList = document.querySelectorAll('[data-bs-toggle="tooltip"]'); + const tooltipList = [...tooltipTriggerList].map(tooltipTriggerEl => new bootstrap.Tooltip(tooltipTriggerEl)); + + loadModels();""" + +new_init = """ // Initialize components + document.addEventListener("DOMContentLoaded", async () => { + // Initialize tooltips + const tooltipTriggerList = document.querySelectorAll('[data-bs-toggle="tooltip"]'); + const tooltipList = [...tooltipTriggerList].map(tooltipTriggerEl => new bootstrap.Tooltip(tooltipTriggerEl)); + + // Load config limits + try { + const res = await fetch('/info'); + if (res.ok) { + const info = await res.json(); + if (info.constraints) { + const stepsEl = document.getElementById('steps'); + if (stepsEl && info.constraints.max_steps) { + stepsEl.max = info.constraints.max_steps; + } + + const widthEl = document.getElementById('width'); + if (widthEl && info.constraints.max_width) { + widthEl.max = info.constraints.max_width; + } + + const heightEl = document.getElementById('height'); + if (heightEl && info.constraints.max_height) { + heightEl.max = info.constraints.max_height; + } + } + } + } catch (e) { + console.error("Failed to load constraints:", e); + } + + loadModels();""" + +content = content.replace(old_init, new_init) + +with open('src/zimage/static/js/main.js', 'w') as f: + f.write(content) diff --git a/patch_js_limits_fix.py b/patch_js_limits_fix.py new file mode 100644 index 0000000..98ad5a3 --- /dev/null +++ b/patch_js_limits_fix.py @@ -0,0 +1,52 @@ +import re + +with open('src/zimage/static/js/main.js', 'r') as f: + content = f.read() + +old_startup = """ console.log("Z-Image Studio: Running startup load..."); + await Promise.all([ + loadModels(), + loadHistory() + ]); + renderActiveLoras(); """ + +new_startup = """ // --- Config Limits Logic --- + async function loadConfig() { + try { + const res = await fetch('/info'); + if (res.ok) { + const info = await res.json(); + if (info.constraints) { + const stepsEl = document.getElementById('steps'); + if (stepsEl && info.constraints.max_steps) { + stepsEl.max = info.constraints.max_steps; + } + + const widthEl = document.getElementById('width'); + if (widthEl && info.constraints.max_width) { + widthEl.max = info.constraints.max_width; + } + + const heightEl = document.getElementById('height'); + if (heightEl && info.constraints.max_height) { + heightEl.max = info.constraints.max_height; + } + } + } + } catch (e) { + console.error("Failed to load constraints:", e); + } + } + + console.log("Z-Image Studio: Running startup load..."); + await Promise.all([ + loadConfig(), + loadModels(), + loadHistory() + ]); + renderActiveLoras(); """ + +content = content.replace(old_startup, new_startup) + +with open('src/zimage/static/js/main.js', 'w') as f: + f.write(content) diff --git a/patch_mcp_limits.py b/patch_mcp_limits.py new file mode 100644 index 0000000..58340c7 --- /dev/null +++ b/patch_mcp_limits.py @@ -0,0 +1,58 @@ +with open('src/zimage/mcp_server.py', 'r') as f: + content = f.read() + +# Let's add load_config to imports +old_imports = """try: + from .hardware import get_available_models, normalize_precision, MODEL_ID_MAP + from . import db + from .storage import save_image, record_generation + from .logger import get_logger, setup_logging +except ImportError: + from hardware import get_available_models, normalize_precision, MODEL_ID_MAP + import db + from storage import save_image, record_generation + from logger import get_logger, setup_logging""" + +new_imports = """try: + from .hardware import get_available_models, normalize_precision, MODEL_ID_MAP + from . import db + from .storage import save_image, record_generation + from .logger import get_logger, setup_logging + from .paths import load_config +except ImportError: + from hardware import get_available_models, normalize_precision, MODEL_ID_MAP + import db + from storage import save_image, record_generation + from logger import get_logger, setup_logging + from paths import load_config""" + +content = content.replace(old_imports, new_imports) + +# Add limits check in _generate_impl +old_impl_start = """ try: + await send_progress(0, "Initializing generation...") + + # Normalize and validate precision""" + +new_impl_start = """ try: + await send_progress(0, "Initializing generation...") + + # Enforce constraints + config = load_config() + max_steps = config.get("max_steps", 50) + max_width = config.get("max_width", 4096) + max_height = config.get("max_height", 4096) + + if steps > max_steps: + raise ValueError(f"Requested steps ({steps}) exceeds the maximum allowed ({max_steps}).") + if width > max_width: + raise ValueError(f"Requested width ({width}) exceeds the maximum allowed ({max_width}).") + if height > max_height: + raise ValueError(f"Requested height ({height}) exceeds the maximum allowed ({max_height}).") + + # Normalize and validate precision""" + +content = content.replace(old_impl_start, new_impl_start) + +with open('src/zimage/mcp_server.py', 'w') as f: + f.write(content) diff --git a/patch_mcp_limits_schema.py b/patch_mcp_limits_schema.py new file mode 100644 index 0000000..7f33e38 --- /dev/null +++ b/patch_mcp_limits_schema.py @@ -0,0 +1,26 @@ +with open('src/zimage/mcp_server.py', 'r') as f: + content = f.read() + +old_func = """@mcp.tool() +async def generate( + prompt: str, + steps: int = 9, + width: int = 1280, + height: int = 720,""" + +# Import Field from pydantic to annotate bounds. FastMCP uses Pydantic. +old_imports = "from urllib.parse import quote" +new_imports = "from urllib.parse import quote\nfrom pydantic import Field" + +new_func = """@mcp.tool() +async def generate( + prompt: str, + steps: int = Field(default=9, description="Number of inference steps (max bounded by server config)"), + width: int = Field(default=1280, description="Image width in pixels (max bounded by server config)"), + height: int = Field(default=720, description="Image height in pixels (max bounded by server config)"),""" + +content = content.replace(old_imports, new_imports) +content = content.replace(old_func, new_func) + +with open('src/zimage/mcp_server.py', 'w') as f: + f.write(content) diff --git a/patch_paths.py b/patch_paths.py new file mode 100644 index 0000000..e857f8b --- /dev/null +++ b/patch_paths.py @@ -0,0 +1,27 @@ +import re + +with open('src/zimage/paths.py', 'r') as f: + content = f.read() + +# Add max constraints to config default +old_config = """ config = { + "version": 1, + "Z_IMAGE_STUDIO_DATA_DIR": None, + "Z_IMAGE_STUDIO_OUTPUT_DIR": None, + "ZIMAGE_ENABLE_TORCH_COMPILE": None, + }""" + +new_config = """ config = { + "version": 1, + "Z_IMAGE_STUDIO_DATA_DIR": None, + "Z_IMAGE_STUDIO_OUTPUT_DIR": None, + "ZIMAGE_ENABLE_TORCH_COMPILE": None, + "max_steps": 50, + "max_width": 4096, + "max_height": 4096, + }""" + +content = content.replace(old_config, new_config) + +with open('src/zimage/paths.py', 'w') as f: + f.write(content) diff --git a/patch_server_info.py b/patch_server_info.py new file mode 100644 index 0000000..f1e4d72 --- /dev/null +++ b/patch_server_info.py @@ -0,0 +1,27 @@ +import re + +with open('src/zimage/server.py', 'r') as f: + content = f.read() + +# I need to add an /api/info endpoint in server.py +old_endpoint = """@app.get("/api/models") +async def get_models():""" + +new_endpoint = """@app.get("/api/info") +async def get_info(): + \"\"\"Return system info including paths, config and hardware constraints.\"\"\" + try: + from .cli import collect_info + return collect_info() + except ImportError: + # Fallback if relative import fails + from cli import collect_info + return collect_info() + +@app.get("/api/models") +async def get_models():""" + +content = content.replace(old_endpoint, new_endpoint) + +with open('src/zimage/server.py', 'w') as f: + f.write(content) diff --git a/patch_server_info_fix.py b/patch_server_info_fix.py new file mode 100644 index 0000000..42d81d6 --- /dev/null +++ b/patch_server_info_fix.py @@ -0,0 +1,30 @@ +import re + +with open('src/zimage/server.py', 'r') as f: + content = f.read() + +# Try again to patch +old_endpoint = """@app.get("/api/models") +async def get_models():""" + +new_endpoint = """@app.get("/api/info") +async def get_info(): + \"\"\"Return system info including paths, config and hardware constraints.\"\"\" + try: + from .cli import collect_info + return collect_info() + except ImportError: + # Fallback if relative import fails + from cli import collect_info + return collect_info() + +@app.get("/api/models") +async def get_models():""" + +if old_endpoint in content: + content = content.replace(old_endpoint, new_endpoint) +else: + print("Could not find endpoint to replace.") + +with open('src/zimage/server.py', 'w') as f: + f.write(content) diff --git a/patch_server_info_fix2.py b/patch_server_info_fix2.py new file mode 100644 index 0000000..09ec7d9 --- /dev/null +++ b/patch_server_info_fix2.py @@ -0,0 +1,27 @@ +with open('src/zimage/server.py', 'r') as f: + content = f.read() + +old_endpoint = """@app.get("/models") +async def get_models():""" + +new_endpoint = """@app.get("/info") +async def get_info(): + \"\"\"Return system info including paths, config and hardware constraints.\"\"\" + try: + from .cli import collect_info + return collect_info() + except ImportError: + # Fallback if relative import fails + import sys + from pathlib import Path + sys.path.insert(0, str(Path(__file__).parent)) + from cli import collect_info + return collect_info() + +@app.get("/models") +async def get_models():""" + +content = content.replace(old_endpoint, new_endpoint) + +with open('src/zimage/server.py', 'w') as f: + f.write(content) diff --git a/patch_server_limits.py b/patch_server_limits.py new file mode 100644 index 0000000..0c96e88 --- /dev/null +++ b/patch_server_limits.py @@ -0,0 +1,31 @@ +with open('src/zimage/server.py', 'r') as f: + content = f.read() + +# Add config check to generate endpoint +old_generate = """@app.post("/generate", response_model=GenerateResponse) +async def generate(req: GenerateRequest, background_tasks: BackgroundTasks): + try: + # Normalize and validate precision early to avoid KeyError inside engine""" + +new_generate = """@app.post("/generate", response_model=GenerateResponse) +async def generate(req: GenerateRequest, background_tasks: BackgroundTasks): + try: + from .paths import load_config + config = load_config() + max_steps = config.get("max_steps", 50) + max_width = config.get("max_width", 4096) + max_height = config.get("max_height", 4096) + + if req.steps > max_steps: + raise HTTPException(status_code=400, detail=f"Requested steps ({req.steps}) exceeds the maximum allowed ({max_steps}).") + if req.width > max_width: + raise HTTPException(status_code=400, detail=f"Requested width ({req.width}) exceeds the maximum allowed ({max_width}).") + if req.height > max_height: + raise HTTPException(status_code=400, detail=f"Requested height ({req.height}) exceeds the maximum allowed ({max_height}).") + + # Normalize and validate precision early to avoid KeyError inside engine""" + +content = content.replace(old_generate, new_generate) + +with open('src/zimage/server.py', 'w') as f: + f.write(content) diff --git a/patch_test_cli_info.py b/patch_test_cli_info.py new file mode 100644 index 0000000..5eeea03 --- /dev/null +++ b/patch_test_cli_info.py @@ -0,0 +1,32 @@ +import re + +with open('tests/test_cli_info.py', 'r') as f: + content = f.read() + +old_paths = """ "paths": { + "module_file": "a", + "config_path": "b", + "data_dir": "c", + "outputs_dir": "d", + "loras_dir": "e", + "db_path": "f", + },""" + +new_paths = """ "paths": { + "module_file": "a", + "config_path": "b", + "data_dir": "c", + "outputs_dir": "d", + "loras_dir": "e", + "db_path": "f", + }, + "constraints": { + "max_steps": 50, + "max_width": 4096, + "max_height": 4096, + },""" + +content = content.replace(old_paths, new_paths) + +with open('tests/test_cli_info.py', 'w') as f: + f.write(content) diff --git a/src/zimage/cli.py b/src/zimage/cli.py index 12d703b..1b14565 100644 --- a/src/zimage/cli.py +++ b/src/zimage/cli.py @@ -27,6 +27,7 @@ get_outputs_dir, get_db_path, get_config_path, + load_config, ) from zimage.logger import get_logger, setup_logging except ImportError: @@ -40,6 +41,7 @@ get_outputs_dir, get_db_path, get_config_path, + load_config, ) from logger import get_logger, setup_logging elif __package__: @@ -52,6 +54,7 @@ get_outputs_dir, get_db_path, get_config_path, + load_config, ) from .logger import get_logger, setup_logging else: @@ -66,6 +69,7 @@ get_outputs_dir, get_db_path, get_config_path, + load_config, ) from logger import get_logger, setup_logging except ImportError: @@ -80,6 +84,7 @@ get_outputs_dir, get_db_path, get_config_path, + load_config, ) from logger import get_logger, setup_logging @@ -304,6 +309,11 @@ def collect_info(): "loras_dir": str(get_loras_dir().resolve()), "db_path": str(get_db_path().resolve()), }, + "constraints": { + "max_steps": load_config().get("max_steps", 50), + "max_width": load_config().get("max_width", 4096), + "max_height": load_config().get("max_height", 4096), + }, "env_overrides": { "Z_IMAGE_STUDIO_DATA_DIR": os.environ.get("Z_IMAGE_STUDIO_DATA_DIR"), "Z_IMAGE_STUDIO_OUTPUT_DIR": os.environ.get("Z_IMAGE_STUDIO_OUTPUT_DIR"), @@ -336,6 +346,11 @@ def format_info_text(info: dict) -> str: f" LoRAs Dir: {info['paths']['loras_dir']}", f" DB Path: {info['paths']['db_path']}", "", + "Constraints:", + f" Max Steps: {info['constraints']['max_steps']}", + f" Max Width: {info['constraints']['max_width']}", + f" Max Height: {info['constraints']['max_height']}", + "", "Environment Overrides:", f" Z_IMAGE_STUDIO_DATA_DIR: {info['env_overrides']['Z_IMAGE_STUDIO_DATA_DIR']}", f" Z_IMAGE_STUDIO_OUTPUT_DIR: {info['env_overrides']['Z_IMAGE_STUDIO_OUTPUT_DIR']}", @@ -421,6 +436,21 @@ def run_generation(args): generate_image, save_image, record_generation = _load_generation_modules() logger.info(f"DEBUG: cwd: {Path.cwd().resolve()}") + # Check constraints + config = load_config() + max_steps = config.get("max_steps", 50) + max_width = config.get("max_width", 4096) + max_height = config.get("max_height", 4096) + + if args.steps > max_steps: + log_error(f"Requested steps ({args.steps}) exceeds the maximum allowed ({max_steps}).") + sys.exit(1) + if args.width > max_width: + log_error(f"Requested width ({args.width}) exceeds the maximum allowed ({max_width}).") + sys.exit(1) + if args.height > max_height: + log_error(f"Requested height ({args.height}) exceeds the maximum allowed ({max_height}).") + # Ensure width/height are multiples of 16 for name in ["width", "height"]: v = getattr(args, name) diff --git a/src/zimage/mcp_server.py b/src/zimage/mcp_server.py index fcf03a4..ed89d08 100644 --- a/src/zimage/mcp_server.py +++ b/src/zimage/mcp_server.py @@ -9,6 +9,7 @@ import base64 import random from urllib.parse import quote +from pydantic import Field # Lazy import for yarl to avoid dependency issues try: @@ -21,11 +22,13 @@ from . import db from .storage import save_image, record_generation from .logger import get_logger, setup_logging + from .paths import load_config except ImportError: from hardware import get_available_models, normalize_precision, MODEL_ID_MAP import db from storage import save_image, record_generation from logger import get_logger, setup_logging + from paths import load_config # Lazy imports for heavy dependencies def _get_engine(): @@ -115,6 +118,19 @@ async def send_progress(percentage: int, message: str): try: await send_progress(0, "Initializing generation...") + # Enforce constraints + config = load_config() + max_steps = config.get("max_steps", 50) + max_width = config.get("max_width", 4096) + max_height = config.get("max_height", 4096) + + if steps > max_steps: + raise ValueError(f"Requested steps ({steps}) exceeds the maximum allowed ({max_steps}).") + if width > max_width: + raise ValueError(f"Requested width ({width}) exceeds the maximum allowed ({max_width}).") + if height > max_height: + raise ValueError(f"Requested height ({height}) exceeds the maximum allowed ({max_height}).") + # Normalize and validate precision try: precision = normalize_precision(precision) @@ -451,9 +467,9 @@ async def send_progress(percentage: int, message: str): @mcp.tool() async def generate( prompt: str, - steps: int = 9, - width: int = 1280, - height: int = 720, + steps: int = Field(default=9, description="Number of inference steps (max bounded by server config)"), + width: int = Field(default=1280, description="Image width in pixels (max bounded by server config)"), + height: int = Field(default=720, description="Image height in pixels (max bounded by server config)"), seed: int | None = None, precision: str = "q8", ctx: Optional[Context] = None diff --git a/src/zimage/paths.py b/src/zimage/paths.py index cc8c40c..0c49c9c 100644 --- a/src/zimage/paths.py +++ b/src/zimage/paths.py @@ -141,6 +141,9 @@ def ensure_initial_setup(): "Z_IMAGE_STUDIO_DATA_DIR": None, "Z_IMAGE_STUDIO_OUTPUT_DIR": None, "ZIMAGE_ENABLE_TORCH_COMPILE": None, + "max_steps": 50, + "max_width": 4096, + "max_height": 4096, } global _CONFIG_CACHE _CONFIG_CACHE = config diff --git a/src/zimage/server.py b/src/zimage/server.py index cbad275..f0cef80 100644 --- a/src/zimage/server.py +++ b/src/zimage/server.py @@ -217,6 +217,20 @@ class GenerateResponse(BaseModel): model_id: str loras: List[LoraInput] = [] +@app.get("/info") +async def get_info(): + """Return system info including paths, config and hardware constraints.""" + try: + from .cli import collect_info + return collect_info() + except ImportError: + # Fallback if relative import fails + import sys + from pathlib import Path + sys.path.insert(0, str(Path(__file__).parent)) + from cli import collect_info + return collect_info() + @app.get("/models") async def get_models(): """Get list of available models with hardware recommendations.""" @@ -372,6 +386,19 @@ async def delete_lora(lora_id: int): @app.post("/generate", response_model=GenerateResponse) async def generate(req: GenerateRequest, background_tasks: BackgroundTasks): try: + from .paths import load_config + config = load_config() + max_steps = config.get("max_steps", 50) + max_width = config.get("max_width", 4096) + max_height = config.get("max_height", 4096) + + if req.steps > max_steps: + raise HTTPException(status_code=400, detail=f"Requested steps ({req.steps}) exceeds the maximum allowed ({max_steps}).") + if req.width > max_width: + raise HTTPException(status_code=400, detail=f"Requested width ({req.width}) exceeds the maximum allowed ({max_width}).") + if req.height > max_height: + raise HTTPException(status_code=400, detail=f"Requested height ({req.height}) exceeds the maximum allowed ({max_height}).") + # Normalize and validate precision early to avoid KeyError inside engine try: precision = normalize_precision(req.precision) diff --git a/src/zimage/static/js/main.js b/src/zimage/static/js/main.js index 1b45c63..c1eef6b 100644 --- a/src/zimage/static/js/main.js +++ b/src/zimage/static/js/main.js @@ -1974,8 +1974,37 @@ async function deleteHistoryItem(itemId) { }; } + // --- Config Limits Logic --- + async function loadConfig() { + try { + const res = await fetch('/info'); + if (res.ok) { + const info = await res.json(); + if (info.constraints) { + const stepsEl = document.getElementById('steps'); + if (stepsEl && info.constraints.max_steps) { + stepsEl.max = info.constraints.max_steps; + } + + const widthEl = document.getElementById('width'); + if (widthEl && info.constraints.max_width) { + widthEl.max = info.constraints.max_width; + } + + const heightEl = document.getElementById('height'); + if (heightEl && info.constraints.max_height) { + heightEl.max = info.constraints.max_height; + } + } + } + } catch (e) { + console.error("Failed to load constraints:", e); + } + } + console.log("Z-Image Studio: Running startup load..."); await Promise.all([ + loadConfig(), loadModels(), loadHistory() ]); diff --git a/test_paths.py b/test_paths.py new file mode 100644 index 0000000..3bb5cc3 --- /dev/null +++ b/test_paths.py @@ -0,0 +1,16 @@ +import sys +from pathlib import Path +sys.path.insert(0, str(Path("src").resolve())) +from zimage.paths import load_config, ensure_initial_setup + +# Mock the config path to test writing +import zimage.paths +import tempfile +temp_dir = tempfile.mkdtemp() +zimage.paths.CONFIG_PATH = Path(temp_dir) / "config.json" +zimage.paths._CONFIG_CACHE = None + +ensure_initial_setup() +config = load_config() +print("Config generated:") +print(config) diff --git a/test_server_info.py b/test_server_info.py new file mode 100644 index 0000000..95980fd --- /dev/null +++ b/test_server_info.py @@ -0,0 +1,12 @@ +import sys +from pathlib import Path +sys.path.insert(0, str(Path("src").resolve())) +import asyncio +from zimage.server import get_info + +async def test(): + info = await get_info() + print("API Info returned keys:", info.keys()) + print("Constraints:", info['constraints']) + +asyncio.run(test()) diff --git a/tests/test_cli_info.py b/tests/test_cli_info.py index e061b31..76f8376 100644 --- a/tests/test_cli_info.py +++ b/tests/test_cli_info.py @@ -71,6 +71,11 @@ def test_run_info_json_outputs_valid_json(capsys, fake_hardware): "loras_dir": "e", "db_path": "f", }, + "constraints": { + "max_steps": 50, + "max_width": 4096, + "max_height": 4096, + }, "env_overrides": { "Z_IMAGE_STUDIO_DATA_DIR": None, "Z_IMAGE_STUDIO_OUTPUT_DIR": None, @@ -108,6 +113,11 @@ def test_run_info_text_includes_hardware_error(capsys): "loras_dir": "e", "db_path": "f", }, + "constraints": { + "max_steps": 50, + "max_width": 4096, + "max_height": 4096, + }, "env_overrides": { "Z_IMAGE_STUDIO_DATA_DIR": None, "Z_IMAGE_STUDIO_OUTPUT_DIR": None, From c93e304beb1eedbd8ab6ab48226c6bf16b423dd8 Mon Sep 17 00:00:00 2001 From: iconben <9401905+iconben@users.noreply.github.com> Date: Sun, 5 Apr 2026 05:23:04 +0000 Subject: [PATCH 2/3] feat: enforce max constraints on image generation limits Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> --- patch_js_models.py | 87 ++++++++++++++++++++++++++++++++++++ patch_server_models.py | 48 ++++++++++++++++++++ src/zimage/server.py | 28 +++++------- src/zimage/static/js/main.js | 46 +++++++------------ test_server_info.py | 12 ----- 5 files changed, 164 insertions(+), 57 deletions(-) create mode 100644 patch_js_models.py create mode 100644 patch_server_models.py delete mode 100644 test_server_info.py diff --git a/patch_js_models.py b/patch_js_models.py new file mode 100644 index 0000000..4e28800 --- /dev/null +++ b/patch_js_models.py @@ -0,0 +1,87 @@ +with open('src/zimage/static/js/main.js', 'r') as f: + content = f.read() + +# 1. Remove the standalone loadConfig call +old_startup = """ // --- Config Limits Logic --- + async function loadConfig() { + try { + const res = await fetch('/info'); + if (res.ok) { + const info = await res.json(); + if (info.constraints) { + const stepsEl = document.getElementById('steps'); + if (stepsEl && info.constraints.max_steps) { + stepsEl.max = info.constraints.max_steps; + } + + const widthEl = document.getElementById('width'); + if (widthEl && info.constraints.max_width) { + widthEl.max = info.constraints.max_width; + } + + const heightEl = document.getElementById('height'); + if (heightEl && info.constraints.max_height) { + heightEl.max = info.constraints.max_height; + } + } + } + } catch (e) { + console.error("Failed to load constraints:", e); + } + } + + console.log("Z-Image Studio: Running startup load..."); + await Promise.all([ + loadConfig(), + loadModels(), + loadHistory() + ]);""" + +new_startup = """ console.log("Z-Image Studio: Running startup load..."); + await Promise.all([ + loadModels(), + loadHistory() + ]);""" + +content = content.replace(old_startup, new_startup) + +# 2. Add limits logic to loadModels +old_load_models = """ // --- Models Loading Logic --- + async function loadModels() { + try { + const res = await fetch('/models'); + const data = await res.json(); + + if (data.device) window.currentDevice = data.device; + if (data.default_precision) window.defaultPrecision = data.default_precision;""" + +new_load_models = """ // --- Models Loading Logic --- + async function loadModels() { + try { + const res = await fetch('/models'); + const data = await res.json(); + + if (data.constraints) { + const stepsEl = document.getElementById('steps'); + if (stepsEl && data.constraints.max_steps) { + stepsEl.max = data.constraints.max_steps; + } + + const widthEl = document.getElementById('width'); + if (widthEl && data.constraints.max_width) { + widthEl.max = data.constraints.max_width; + } + + const heightEl = document.getElementById('height'); + if (heightEl && data.constraints.max_height) { + heightEl.max = data.constraints.max_height; + } + } + + if (data.device) window.currentDevice = data.device; + if (data.default_precision) window.defaultPrecision = data.default_precision;""" + +content = content.replace(old_load_models, new_load_models) + +with open('src/zimage/static/js/main.js', 'w') as f: + f.write(content) diff --git a/patch_server_models.py b/patch_server_models.py new file mode 100644 index 0000000..d30741c --- /dev/null +++ b/patch_server_models.py @@ -0,0 +1,48 @@ +import re + +with open('src/zimage/server.py', 'r') as f: + content = f.read() + +# 1. Remove the /info endpoint we added +old_info = """@app.get("/info") +async def get_info(): + \"\"\"Return system info including paths, config and hardware constraints.\"\"\" + try: + from .cli import collect_info + return collect_info() + except ImportError: + # Fallback if relative import fails + import sys + from pathlib import Path + sys.path.insert(0, str(Path(__file__).parent)) + from cli import collect_info + return collect_info() + +""" +content = content.replace(old_info, "") + +# 2. Update /models endpoint to include constraints +old_models = """@app.get("/models") +async def get_models(): + \"\"\"Get list of available models with hardware recommendations.\"\"\" + return get_available_models()""" + +new_models = """@app.get("/models") +async def get_models(): + \"\"\"Get list of available models with hardware recommendations and constraints.\"\"\" + from .paths import load_config + config = load_config() + models_info = get_available_models() + + # Inject constraints into the response + models_info['constraints'] = { + "max_steps": config.get("max_steps", 50), + "max_width": config.get("max_width", 4096), + "max_height": config.get("max_height", 4096), + } + return models_info""" + +content = content.replace(old_models, new_models) + +with open('src/zimage/server.py', 'w') as f: + f.write(content) diff --git a/src/zimage/server.py b/src/zimage/server.py index f0cef80..503aa54 100644 --- a/src/zimage/server.py +++ b/src/zimage/server.py @@ -217,24 +217,20 @@ class GenerateResponse(BaseModel): model_id: str loras: List[LoraInput] = [] -@app.get("/info") -async def get_info(): - """Return system info including paths, config and hardware constraints.""" - try: - from .cli import collect_info - return collect_info() - except ImportError: - # Fallback if relative import fails - import sys - from pathlib import Path - sys.path.insert(0, str(Path(__file__).parent)) - from cli import collect_info - return collect_info() - @app.get("/models") async def get_models(): - """Get list of available models with hardware recommendations.""" - return get_available_models() + """Get list of available models with hardware recommendations and constraints.""" + from .paths import load_config + config = load_config() + models_info = get_available_models() + + # Inject constraints into the response + models_info['constraints'] = { + "max_steps": config.get("max_steps", 50), + "max_width": config.get("max_width", 4096), + "max_height": config.get("max_height", 4096), + } + return models_info @app.get("/loras") async def get_loras(): diff --git a/src/zimage/static/js/main.js b/src/zimage/static/js/main.js index c1eef6b..bfece9b 100644 --- a/src/zimage/static/js/main.js +++ b/src/zimage/static/js/main.js @@ -895,6 +895,23 @@ const res = await fetch('/models'); const data = await res.json(); + if (data.constraints) { + const stepsEl = document.getElementById('steps'); + if (stepsEl && data.constraints.max_steps) { + stepsEl.max = data.constraints.max_steps; + } + + const widthEl = document.getElementById('width'); + if (widthEl && data.constraints.max_width) { + widthEl.max = data.constraints.max_width; + } + + const heightEl = document.getElementById('height'); + if (heightEl && data.constraints.max_height) { + heightEl.max = data.constraints.max_height; + } + } + if (data.device) window.currentDevice = data.device; if (data.default_precision) window.defaultPrecision = data.default_precision; @@ -1974,37 +1991,8 @@ async function deleteHistoryItem(itemId) { }; } - // --- Config Limits Logic --- - async function loadConfig() { - try { - const res = await fetch('/info'); - if (res.ok) { - const info = await res.json(); - if (info.constraints) { - const stepsEl = document.getElementById('steps'); - if (stepsEl && info.constraints.max_steps) { - stepsEl.max = info.constraints.max_steps; - } - - const widthEl = document.getElementById('width'); - if (widthEl && info.constraints.max_width) { - widthEl.max = info.constraints.max_width; - } - - const heightEl = document.getElementById('height'); - if (heightEl && info.constraints.max_height) { - heightEl.max = info.constraints.max_height; - } - } - } - } catch (e) { - console.error("Failed to load constraints:", e); - } - } - console.log("Z-Image Studio: Running startup load..."); await Promise.all([ - loadConfig(), loadModels(), loadHistory() ]); diff --git a/test_server_info.py b/test_server_info.py deleted file mode 100644 index 95980fd..0000000 --- a/test_server_info.py +++ /dev/null @@ -1,12 +0,0 @@ -import sys -from pathlib import Path -sys.path.insert(0, str(Path("src").resolve())) -import asyncio -from zimage.server import get_info - -async def test(): - info = await get_info() - print("API Info returned keys:", info.keys()) - print("Constraints:", info['constraints']) - -asyncio.run(test()) From 58f9846112ccbede32f8fe7e731092b278e7287c Mon Sep 17 00:00:00 2001 From: iconben <9401905+iconben@users.noreply.github.com> Date: Sun, 5 Apr 2026 05:53:15 +0000 Subject: [PATCH 3/3] feat: enforce max constraints on image generation limits Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> --- patch_cli_info.py | 58 ------------------------- patch_cli_info_fix.py | 11 ----- patch_cli_info_fix2.py | 32 -------------- patch_cli_info_fix3.py | 26 ----------- patch_cli_info_fix4.py | 34 --------------- patch_cli_info_imports.py | 27 ------------ patch_cli_info_imports_2.py | 11 ----- patch_cli_limits.py | 34 --------------- patch_js_limits.py | 50 --------------------- patch_js_limits_fix.py | 52 ---------------------- patch_js_models.py | 87 ------------------------------------- patch_mcp_limits.py | 58 ------------------------- patch_mcp_limits_schema.py | 26 ----------- patch_paths.py | 27 ------------ patch_server_info.py | 27 ------------ patch_server_info_fix.py | 30 ------------- patch_server_info_fix2.py | 27 ------------ patch_server_limits.py | 31 ------------- patch_server_models.py | 48 -------------------- patch_test_cli_info.py | 32 -------------- test_paths.py | 16 ------- 21 files changed, 744 deletions(-) delete mode 100644 patch_cli_info.py delete mode 100644 patch_cli_info_fix.py delete mode 100644 patch_cli_info_fix2.py delete mode 100644 patch_cli_info_fix3.py delete mode 100644 patch_cli_info_fix4.py delete mode 100644 patch_cli_info_imports.py delete mode 100644 patch_cli_info_imports_2.py delete mode 100644 patch_cli_limits.py delete mode 100644 patch_js_limits.py delete mode 100644 patch_js_limits_fix.py delete mode 100644 patch_js_models.py delete mode 100644 patch_mcp_limits.py delete mode 100644 patch_mcp_limits_schema.py delete mode 100644 patch_paths.py delete mode 100644 patch_server_info.py delete mode 100644 patch_server_info_fix.py delete mode 100644 patch_server_info_fix2.py delete mode 100644 patch_server_limits.py delete mode 100644 patch_server_models.py delete mode 100644 patch_test_cli_info.py delete mode 100644 test_paths.py diff --git a/patch_cli_info.py b/patch_cli_info.py deleted file mode 100644 index e4c3897..0000000 --- a/patch_cli_info.py +++ /dev/null @@ -1,58 +0,0 @@ -import re - -with open('src/zimage/cli.py', 'r') as f: - content = f.read() - -# 1. Update collect_info -old_collect = """ "paths": { - "module_file": str(Path(__file__).resolve()), - "config_path": str(get_config_path().resolve()), - "data_dir": str(get_data_dir().resolve()), - "outputs_dir": str(get_outputs_dir().resolve()), - "loras_dir": str(get_loras_dir().resolve()), - "db_path": str(get_db_path().resolve()), - }, - "hardware": hardware, - }""" - -new_collect = """ "paths": { - "module_file": str(Path(__file__).resolve()), - "config_path": str(get_config_path().resolve()), - "data_dir": str(get_data_dir().resolve()), - "outputs_dir": str(get_outputs_dir().resolve()), - "loras_dir": str(get_loras_dir().resolve()), - "db_path": str(get_db_path().resolve()), - }, - "constraints": { - "max_steps": load_config().get("max_steps", 50), - "max_width": load_config().get("max_width", 4096), - "max_height": load_config().get("max_height", 4096), - }, - "hardware": hardware, - }""" -content = content.replace(old_collect, new_collect) - - -# 2. Update format_info_text -old_format = """ f" Outputs Dir: {info['paths']['outputs_dir']}", - f" LoRAs Dir: {info['paths']['loras_dir']}", - f" DB Path: {info['paths']['db_path']}", - "", - "Hardware:" - ]""" - -new_format = """ f" Outputs Dir: {info['paths']['outputs_dir']}", - f" LoRAs Dir: {info['paths']['loras_dir']}", - f" DB Path: {info['paths']['db_path']}", - "", - "Constraints:", - f" Max Steps: {info['constraints']['max_steps']}", - f" Max Width: {info['constraints']['max_width']}", - f" Max Height: {info['constraints']['max_height']}", - "", - "Hardware:" - ]""" -content = content.replace(old_format, new_format) - -with open('src/zimage/cli.py', 'w') as f: - f.write(content) diff --git a/patch_cli_info_fix.py b/patch_cli_info_fix.py deleted file mode 100644 index 8551c38..0000000 --- a/patch_cli_info_fix.py +++ /dev/null @@ -1,11 +0,0 @@ -import re - -with open('src/zimage/cli.py', 'r') as f: - content = f.read() - -# Make sure constraints are actually output (I noticed they didn't show up in the output) -# Let's check where they should be inserted. -# Ah, I replaced "Environment Overrides:" with "Hardware:", I might have lost "Environment Overrides:". - -with open('src/zimage/cli.py', 'w') as f: - f.write(content) diff --git a/patch_cli_info_fix2.py b/patch_cli_info_fix2.py deleted file mode 100644 index a1bdd7c..0000000 --- a/patch_cli_info_fix2.py +++ /dev/null @@ -1,32 +0,0 @@ -import re - -with open('src/zimage/cli.py', 'r') as f: - content = f.read() - -old_str = """ f" Outputs Dir: {info['paths']['outputs_dir']}", - f" LoRAs Dir: {info['paths']['loras_dir']}", - f" DB Path: {info['paths']['db_path']}", - "", - "Environment Overrides:" - ]""" - -new_str = """ f" Outputs Dir: {info['paths']['outputs_dir']}", - f" LoRAs Dir: {info['paths']['loras_dir']}", - f" DB Path: {info['paths']['db_path']}", - "", - "Constraints:", - f" Max Steps: {info['constraints']['max_steps']}", - f" Max Width: {info['constraints']['max_width']}", - f" Max Height: {info['constraints']['max_height']}", - "", - "Environment Overrides:" - ]""" - -if old_str in content: - content = content.replace(old_str, new_str) -else: - # My previous replacement might have matched something else or failed - pass - -with open('src/zimage/cli.py', 'w') as f: - f.write(content) diff --git a/patch_cli_info_fix3.py b/patch_cli_info_fix3.py deleted file mode 100644 index 8cd8e4a..0000000 --- a/patch_cli_info_fix3.py +++ /dev/null @@ -1,26 +0,0 @@ -import re - -with open('src/zimage/cli.py', 'r') as f: - content = f.read() - -old_str = """ f" Outputs Dir: {info['paths']['outputs_dir']}", - f" LoRAs Dir: {info['paths']['loras_dir']}", - f" DB Path: {info['paths']['db_path']}", - "", - "Environment Overrides:""" - -new_str = """ f" Outputs Dir: {info['paths']['outputs_dir']}", - f" LoRAs Dir: {info['paths']['loras_dir']}", - f" DB Path: {info['paths']['db_path']}", - "", - "Constraints:", - f" Max Steps: {info['constraints']['max_steps']}", - f" Max Width: {info['constraints']['max_width']}", - f" Max Height: {info['constraints']['max_height']}", - "", - "Environment Overrides:""" - -content = content.replace(old_str, new_str) - -with open('src/zimage/cli.py', 'w') as f: - f.write(content) diff --git a/patch_cli_info_fix4.py b/patch_cli_info_fix4.py deleted file mode 100644 index 0f88dee..0000000 --- a/patch_cli_info_fix4.py +++ /dev/null @@ -1,34 +0,0 @@ -with open('src/zimage/cli.py', 'r') as f: - content = f.read() - -# Let's see what collect_info returns currently -old_collect = """ "paths": { - "module_file": str(Path(__file__).resolve()), - "config_path": str(get_config_path().resolve()), - "data_dir": str(get_data_dir().resolve()), - "outputs_dir": str(get_outputs_dir().resolve()), - "loras_dir": str(get_loras_dir().resolve()), - "db_path": str(get_db_path().resolve()), - }, - "env_overrides": {""" - -new_collect = """ "paths": { - "module_file": str(Path(__file__).resolve()), - "config_path": str(get_config_path().resolve()), - "data_dir": str(get_data_dir().resolve()), - "outputs_dir": str(get_outputs_dir().resolve()), - "loras_dir": str(get_loras_dir().resolve()), - "db_path": str(get_db_path().resolve()), - }, - "constraints": { - "max_steps": load_config().get("max_steps", 50), - "max_width": load_config().get("max_width", 4096), - "max_height": load_config().get("max_height", 4096), - }, - "env_overrides": {""" - -if old_collect in content: - content = content.replace(old_collect, new_collect) - -with open('src/zimage/cli.py', 'w') as f: - f.write(content) diff --git a/patch_cli_info_imports.py b/patch_cli_info_imports.py deleted file mode 100644 index fa61d93..0000000 --- a/patch_cli_info_imports.py +++ /dev/null @@ -1,27 +0,0 @@ -import re - -with open('src/zimage/cli.py', 'r') as f: - content = f.read() - -# I need to add load_config to the imports from paths -content = content.replace( - """get_outputs_dir, - get_loras_dir, - get_db_path,""", - """get_outputs_dir, - get_loras_dir, - get_db_path, - load_config,""" -) -content = content.replace( - """get_outputs_dir, - get_loras_dir, - get_db_path,""", - """get_outputs_dir, - get_loras_dir, - get_db_path, - load_config,""" -) - -with open('src/zimage/cli.py', 'w') as f: - f.write(content) diff --git a/patch_cli_info_imports_2.py b/patch_cli_info_imports_2.py deleted file mode 100644 index de42608..0000000 --- a/patch_cli_info_imports_2.py +++ /dev/null @@ -1,11 +0,0 @@ -import re - -with open('src/zimage/cli.py', 'r') as f: - content = f.read() - -content = content.replace("get_db_path,\n get_config_path,\n", "get_db_path,\n get_config_path,\n load_config,\n") -content = content.replace("get_db_path,\n get_config_path,\n", "get_db_path,\n get_config_path,\n load_config,\n") -content = content.replace("get_db_path,\n get_config_path,\n", "get_db_path,\n get_config_path,\n load_config,\n") - -with open('src/zimage/cli.py', 'w') as f: - f.write(content) diff --git a/patch_cli_limits.py b/patch_cli_limits.py deleted file mode 100644 index c5ca3e7..0000000 --- a/patch_cli_limits.py +++ /dev/null @@ -1,34 +0,0 @@ -with open('src/zimage/cli.py', 'r') as f: - content = f.read() - -old_run_gen = """def run_generation(args): - generate_image, save_image, record_generation = _load_generation_modules() - logger.info(f"DEBUG: cwd: {Path.cwd().resolve()}") - - # Ensure width/height are multiples of 16""" - -new_run_gen = """def run_generation(args): - generate_image, save_image, record_generation = _load_generation_modules() - logger.info(f"DEBUG: cwd: {Path.cwd().resolve()}") - - # Check constraints - config = load_config() - max_steps = config.get("max_steps", 50) - max_width = config.get("max_width", 4096) - max_height = config.get("max_height", 4096) - - if args.steps > max_steps: - log_error(f"Requested steps ({args.steps}) exceeds the maximum allowed ({max_steps}).") - sys.exit(1) - if args.width > max_width: - log_error(f"Requested width ({args.width}) exceeds the maximum allowed ({max_width}).") - sys.exit(1) - if args.height > max_height: - log_error(f"Requested height ({args.height}) exceeds the maximum allowed ({max_height}).") - - # Ensure width/height are multiples of 16""" - -content = content.replace(old_run_gen, new_run_gen) - -with open('src/zimage/cli.py', 'w') as f: - f.write(content) diff --git a/patch_js_limits.py b/patch_js_limits.py deleted file mode 100644 index 5589832..0000000 --- a/patch_js_limits.py +++ /dev/null @@ -1,50 +0,0 @@ -with open('src/zimage/static/js/main.js', 'r') as f: - content = f.read() - -# Let's insert a call to /info at startup and update DOM inputs -old_init = """ // Initialize components - document.addEventListener("DOMContentLoaded", () => { - // Initialize tooltips - const tooltipTriggerList = document.querySelectorAll('[data-bs-toggle="tooltip"]'); - const tooltipList = [...tooltipTriggerList].map(tooltipTriggerEl => new bootstrap.Tooltip(tooltipTriggerEl)); - - loadModels();""" - -new_init = """ // Initialize components - document.addEventListener("DOMContentLoaded", async () => { - // Initialize tooltips - const tooltipTriggerList = document.querySelectorAll('[data-bs-toggle="tooltip"]'); - const tooltipList = [...tooltipTriggerList].map(tooltipTriggerEl => new bootstrap.Tooltip(tooltipTriggerEl)); - - // Load config limits - try { - const res = await fetch('/info'); - if (res.ok) { - const info = await res.json(); - if (info.constraints) { - const stepsEl = document.getElementById('steps'); - if (stepsEl && info.constraints.max_steps) { - stepsEl.max = info.constraints.max_steps; - } - - const widthEl = document.getElementById('width'); - if (widthEl && info.constraints.max_width) { - widthEl.max = info.constraints.max_width; - } - - const heightEl = document.getElementById('height'); - if (heightEl && info.constraints.max_height) { - heightEl.max = info.constraints.max_height; - } - } - } - } catch (e) { - console.error("Failed to load constraints:", e); - } - - loadModels();""" - -content = content.replace(old_init, new_init) - -with open('src/zimage/static/js/main.js', 'w') as f: - f.write(content) diff --git a/patch_js_limits_fix.py b/patch_js_limits_fix.py deleted file mode 100644 index 98ad5a3..0000000 --- a/patch_js_limits_fix.py +++ /dev/null @@ -1,52 +0,0 @@ -import re - -with open('src/zimage/static/js/main.js', 'r') as f: - content = f.read() - -old_startup = """ console.log("Z-Image Studio: Running startup load..."); - await Promise.all([ - loadModels(), - loadHistory() - ]); - renderActiveLoras(); """ - -new_startup = """ // --- Config Limits Logic --- - async function loadConfig() { - try { - const res = await fetch('/info'); - if (res.ok) { - const info = await res.json(); - if (info.constraints) { - const stepsEl = document.getElementById('steps'); - if (stepsEl && info.constraints.max_steps) { - stepsEl.max = info.constraints.max_steps; - } - - const widthEl = document.getElementById('width'); - if (widthEl && info.constraints.max_width) { - widthEl.max = info.constraints.max_width; - } - - const heightEl = document.getElementById('height'); - if (heightEl && info.constraints.max_height) { - heightEl.max = info.constraints.max_height; - } - } - } - } catch (e) { - console.error("Failed to load constraints:", e); - } - } - - console.log("Z-Image Studio: Running startup load..."); - await Promise.all([ - loadConfig(), - loadModels(), - loadHistory() - ]); - renderActiveLoras(); """ - -content = content.replace(old_startup, new_startup) - -with open('src/zimage/static/js/main.js', 'w') as f: - f.write(content) diff --git a/patch_js_models.py b/patch_js_models.py deleted file mode 100644 index 4e28800..0000000 --- a/patch_js_models.py +++ /dev/null @@ -1,87 +0,0 @@ -with open('src/zimage/static/js/main.js', 'r') as f: - content = f.read() - -# 1. Remove the standalone loadConfig call -old_startup = """ // --- Config Limits Logic --- - async function loadConfig() { - try { - const res = await fetch('/info'); - if (res.ok) { - const info = await res.json(); - if (info.constraints) { - const stepsEl = document.getElementById('steps'); - if (stepsEl && info.constraints.max_steps) { - stepsEl.max = info.constraints.max_steps; - } - - const widthEl = document.getElementById('width'); - if (widthEl && info.constraints.max_width) { - widthEl.max = info.constraints.max_width; - } - - const heightEl = document.getElementById('height'); - if (heightEl && info.constraints.max_height) { - heightEl.max = info.constraints.max_height; - } - } - } - } catch (e) { - console.error("Failed to load constraints:", e); - } - } - - console.log("Z-Image Studio: Running startup load..."); - await Promise.all([ - loadConfig(), - loadModels(), - loadHistory() - ]);""" - -new_startup = """ console.log("Z-Image Studio: Running startup load..."); - await Promise.all([ - loadModels(), - loadHistory() - ]);""" - -content = content.replace(old_startup, new_startup) - -# 2. Add limits logic to loadModels -old_load_models = """ // --- Models Loading Logic --- - async function loadModels() { - try { - const res = await fetch('/models'); - const data = await res.json(); - - if (data.device) window.currentDevice = data.device; - if (data.default_precision) window.defaultPrecision = data.default_precision;""" - -new_load_models = """ // --- Models Loading Logic --- - async function loadModels() { - try { - const res = await fetch('/models'); - const data = await res.json(); - - if (data.constraints) { - const stepsEl = document.getElementById('steps'); - if (stepsEl && data.constraints.max_steps) { - stepsEl.max = data.constraints.max_steps; - } - - const widthEl = document.getElementById('width'); - if (widthEl && data.constraints.max_width) { - widthEl.max = data.constraints.max_width; - } - - const heightEl = document.getElementById('height'); - if (heightEl && data.constraints.max_height) { - heightEl.max = data.constraints.max_height; - } - } - - if (data.device) window.currentDevice = data.device; - if (data.default_precision) window.defaultPrecision = data.default_precision;""" - -content = content.replace(old_load_models, new_load_models) - -with open('src/zimage/static/js/main.js', 'w') as f: - f.write(content) diff --git a/patch_mcp_limits.py b/patch_mcp_limits.py deleted file mode 100644 index 58340c7..0000000 --- a/patch_mcp_limits.py +++ /dev/null @@ -1,58 +0,0 @@ -with open('src/zimage/mcp_server.py', 'r') as f: - content = f.read() - -# Let's add load_config to imports -old_imports = """try: - from .hardware import get_available_models, normalize_precision, MODEL_ID_MAP - from . import db - from .storage import save_image, record_generation - from .logger import get_logger, setup_logging -except ImportError: - from hardware import get_available_models, normalize_precision, MODEL_ID_MAP - import db - from storage import save_image, record_generation - from logger import get_logger, setup_logging""" - -new_imports = """try: - from .hardware import get_available_models, normalize_precision, MODEL_ID_MAP - from . import db - from .storage import save_image, record_generation - from .logger import get_logger, setup_logging - from .paths import load_config -except ImportError: - from hardware import get_available_models, normalize_precision, MODEL_ID_MAP - import db - from storage import save_image, record_generation - from logger import get_logger, setup_logging - from paths import load_config""" - -content = content.replace(old_imports, new_imports) - -# Add limits check in _generate_impl -old_impl_start = """ try: - await send_progress(0, "Initializing generation...") - - # Normalize and validate precision""" - -new_impl_start = """ try: - await send_progress(0, "Initializing generation...") - - # Enforce constraints - config = load_config() - max_steps = config.get("max_steps", 50) - max_width = config.get("max_width", 4096) - max_height = config.get("max_height", 4096) - - if steps > max_steps: - raise ValueError(f"Requested steps ({steps}) exceeds the maximum allowed ({max_steps}).") - if width > max_width: - raise ValueError(f"Requested width ({width}) exceeds the maximum allowed ({max_width}).") - if height > max_height: - raise ValueError(f"Requested height ({height}) exceeds the maximum allowed ({max_height}).") - - # Normalize and validate precision""" - -content = content.replace(old_impl_start, new_impl_start) - -with open('src/zimage/mcp_server.py', 'w') as f: - f.write(content) diff --git a/patch_mcp_limits_schema.py b/patch_mcp_limits_schema.py deleted file mode 100644 index 7f33e38..0000000 --- a/patch_mcp_limits_schema.py +++ /dev/null @@ -1,26 +0,0 @@ -with open('src/zimage/mcp_server.py', 'r') as f: - content = f.read() - -old_func = """@mcp.tool() -async def generate( - prompt: str, - steps: int = 9, - width: int = 1280, - height: int = 720,""" - -# Import Field from pydantic to annotate bounds. FastMCP uses Pydantic. -old_imports = "from urllib.parse import quote" -new_imports = "from urllib.parse import quote\nfrom pydantic import Field" - -new_func = """@mcp.tool() -async def generate( - prompt: str, - steps: int = Field(default=9, description="Number of inference steps (max bounded by server config)"), - width: int = Field(default=1280, description="Image width in pixels (max bounded by server config)"), - height: int = Field(default=720, description="Image height in pixels (max bounded by server config)"),""" - -content = content.replace(old_imports, new_imports) -content = content.replace(old_func, new_func) - -with open('src/zimage/mcp_server.py', 'w') as f: - f.write(content) diff --git a/patch_paths.py b/patch_paths.py deleted file mode 100644 index e857f8b..0000000 --- a/patch_paths.py +++ /dev/null @@ -1,27 +0,0 @@ -import re - -with open('src/zimage/paths.py', 'r') as f: - content = f.read() - -# Add max constraints to config default -old_config = """ config = { - "version": 1, - "Z_IMAGE_STUDIO_DATA_DIR": None, - "Z_IMAGE_STUDIO_OUTPUT_DIR": None, - "ZIMAGE_ENABLE_TORCH_COMPILE": None, - }""" - -new_config = """ config = { - "version": 1, - "Z_IMAGE_STUDIO_DATA_DIR": None, - "Z_IMAGE_STUDIO_OUTPUT_DIR": None, - "ZIMAGE_ENABLE_TORCH_COMPILE": None, - "max_steps": 50, - "max_width": 4096, - "max_height": 4096, - }""" - -content = content.replace(old_config, new_config) - -with open('src/zimage/paths.py', 'w') as f: - f.write(content) diff --git a/patch_server_info.py b/patch_server_info.py deleted file mode 100644 index f1e4d72..0000000 --- a/patch_server_info.py +++ /dev/null @@ -1,27 +0,0 @@ -import re - -with open('src/zimage/server.py', 'r') as f: - content = f.read() - -# I need to add an /api/info endpoint in server.py -old_endpoint = """@app.get("/api/models") -async def get_models():""" - -new_endpoint = """@app.get("/api/info") -async def get_info(): - \"\"\"Return system info including paths, config and hardware constraints.\"\"\" - try: - from .cli import collect_info - return collect_info() - except ImportError: - # Fallback if relative import fails - from cli import collect_info - return collect_info() - -@app.get("/api/models") -async def get_models():""" - -content = content.replace(old_endpoint, new_endpoint) - -with open('src/zimage/server.py', 'w') as f: - f.write(content) diff --git a/patch_server_info_fix.py b/patch_server_info_fix.py deleted file mode 100644 index 42d81d6..0000000 --- a/patch_server_info_fix.py +++ /dev/null @@ -1,30 +0,0 @@ -import re - -with open('src/zimage/server.py', 'r') as f: - content = f.read() - -# Try again to patch -old_endpoint = """@app.get("/api/models") -async def get_models():""" - -new_endpoint = """@app.get("/api/info") -async def get_info(): - \"\"\"Return system info including paths, config and hardware constraints.\"\"\" - try: - from .cli import collect_info - return collect_info() - except ImportError: - # Fallback if relative import fails - from cli import collect_info - return collect_info() - -@app.get("/api/models") -async def get_models():""" - -if old_endpoint in content: - content = content.replace(old_endpoint, new_endpoint) -else: - print("Could not find endpoint to replace.") - -with open('src/zimage/server.py', 'w') as f: - f.write(content) diff --git a/patch_server_info_fix2.py b/patch_server_info_fix2.py deleted file mode 100644 index 09ec7d9..0000000 --- a/patch_server_info_fix2.py +++ /dev/null @@ -1,27 +0,0 @@ -with open('src/zimage/server.py', 'r') as f: - content = f.read() - -old_endpoint = """@app.get("/models") -async def get_models():""" - -new_endpoint = """@app.get("/info") -async def get_info(): - \"\"\"Return system info including paths, config and hardware constraints.\"\"\" - try: - from .cli import collect_info - return collect_info() - except ImportError: - # Fallback if relative import fails - import sys - from pathlib import Path - sys.path.insert(0, str(Path(__file__).parent)) - from cli import collect_info - return collect_info() - -@app.get("/models") -async def get_models():""" - -content = content.replace(old_endpoint, new_endpoint) - -with open('src/zimage/server.py', 'w') as f: - f.write(content) diff --git a/patch_server_limits.py b/patch_server_limits.py deleted file mode 100644 index 0c96e88..0000000 --- a/patch_server_limits.py +++ /dev/null @@ -1,31 +0,0 @@ -with open('src/zimage/server.py', 'r') as f: - content = f.read() - -# Add config check to generate endpoint -old_generate = """@app.post("/generate", response_model=GenerateResponse) -async def generate(req: GenerateRequest, background_tasks: BackgroundTasks): - try: - # Normalize and validate precision early to avoid KeyError inside engine""" - -new_generate = """@app.post("/generate", response_model=GenerateResponse) -async def generate(req: GenerateRequest, background_tasks: BackgroundTasks): - try: - from .paths import load_config - config = load_config() - max_steps = config.get("max_steps", 50) - max_width = config.get("max_width", 4096) - max_height = config.get("max_height", 4096) - - if req.steps > max_steps: - raise HTTPException(status_code=400, detail=f"Requested steps ({req.steps}) exceeds the maximum allowed ({max_steps}).") - if req.width > max_width: - raise HTTPException(status_code=400, detail=f"Requested width ({req.width}) exceeds the maximum allowed ({max_width}).") - if req.height > max_height: - raise HTTPException(status_code=400, detail=f"Requested height ({req.height}) exceeds the maximum allowed ({max_height}).") - - # Normalize and validate precision early to avoid KeyError inside engine""" - -content = content.replace(old_generate, new_generate) - -with open('src/zimage/server.py', 'w') as f: - f.write(content) diff --git a/patch_server_models.py b/patch_server_models.py deleted file mode 100644 index d30741c..0000000 --- a/patch_server_models.py +++ /dev/null @@ -1,48 +0,0 @@ -import re - -with open('src/zimage/server.py', 'r') as f: - content = f.read() - -# 1. Remove the /info endpoint we added -old_info = """@app.get("/info") -async def get_info(): - \"\"\"Return system info including paths, config and hardware constraints.\"\"\" - try: - from .cli import collect_info - return collect_info() - except ImportError: - # Fallback if relative import fails - import sys - from pathlib import Path - sys.path.insert(0, str(Path(__file__).parent)) - from cli import collect_info - return collect_info() - -""" -content = content.replace(old_info, "") - -# 2. Update /models endpoint to include constraints -old_models = """@app.get("/models") -async def get_models(): - \"\"\"Get list of available models with hardware recommendations.\"\"\" - return get_available_models()""" - -new_models = """@app.get("/models") -async def get_models(): - \"\"\"Get list of available models with hardware recommendations and constraints.\"\"\" - from .paths import load_config - config = load_config() - models_info = get_available_models() - - # Inject constraints into the response - models_info['constraints'] = { - "max_steps": config.get("max_steps", 50), - "max_width": config.get("max_width", 4096), - "max_height": config.get("max_height", 4096), - } - return models_info""" - -content = content.replace(old_models, new_models) - -with open('src/zimage/server.py', 'w') as f: - f.write(content) diff --git a/patch_test_cli_info.py b/patch_test_cli_info.py deleted file mode 100644 index 5eeea03..0000000 --- a/patch_test_cli_info.py +++ /dev/null @@ -1,32 +0,0 @@ -import re - -with open('tests/test_cli_info.py', 'r') as f: - content = f.read() - -old_paths = """ "paths": { - "module_file": "a", - "config_path": "b", - "data_dir": "c", - "outputs_dir": "d", - "loras_dir": "e", - "db_path": "f", - },""" - -new_paths = """ "paths": { - "module_file": "a", - "config_path": "b", - "data_dir": "c", - "outputs_dir": "d", - "loras_dir": "e", - "db_path": "f", - }, - "constraints": { - "max_steps": 50, - "max_width": 4096, - "max_height": 4096, - },""" - -content = content.replace(old_paths, new_paths) - -with open('tests/test_cli_info.py', 'w') as f: - f.write(content) diff --git a/test_paths.py b/test_paths.py deleted file mode 100644 index 3bb5cc3..0000000 --- a/test_paths.py +++ /dev/null @@ -1,16 +0,0 @@ -import sys -from pathlib import Path -sys.path.insert(0, str(Path("src").resolve())) -from zimage.paths import load_config, ensure_initial_setup - -# Mock the config path to test writing -import zimage.paths -import tempfile -temp_dir = tempfile.mkdtemp() -zimage.paths.CONFIG_PATH = Path(temp_dir) / "config.json" -zimage.paths._CONFIG_CACHE = None - -ensure_initial_setup() -config = load_config() -print("Config generated:") -print(config)