from __future__ import annotations import base64 import os import platform import tempfile from pathlib import Path from typing import Any, Dict from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from fastapi.staticfiles import StaticFiles from .config import AppConfig, load_config, save_config from .models import ConfigUpdate, I2IRequest, InpaintRequest, T2IRequest from .services.novelai import generate_i2i, generate_inpaint, generate_t2i def _ensure_x64(n: int) -> int: if n <= 64: return 64 if n % 64 == 0: return n return ((n // 64) + 1) * 64 if (n / 64) % 1 >= 0.5 else (n // 64) * 64 def _as_data_uri(b64: str) -> str: return f"data:image/png;base64,{b64}" app = FastAPI(title="New NAI", version="1.0.0") # CORS app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=False, allow_methods=["*"], allow_headers=["*"], ) @app.get("/api/health") def health() -> Dict[str, Any]: return {"status": "ok"} @app.get("/api/config") def get_config() -> Dict[str, Any]: cfg: AppConfig = load_config() return cfg.model_dump() @app.put("/api/config") def update_config(update: ConfigUpdate) -> Dict[str, Any]: cfg = load_config() data = update.model_dump(exclude_none=True) for k, v in data.items(): setattr(cfg, k, v) save_config(cfg) return cfg.model_dump() @app.get("/api/select-output-dir") def api_select_output_dir() -> Dict[str, Any]: """ 在本机弹出目录选择器并返回选择的目录路径。 多重回退方案: 1) tkinter(如可用) 2) Windows: PowerShell + Shell.Application.BrowseForFolder 3) macOS: osascript choose folder 4) Linux: zenity --file-selection --directory 全部失败则返回 500。 """ # 1) tkinter try: import tkinter as tk # type: ignore from tkinter import filedialog # type: ignore root = tk.Tk() root.withdraw() path = filedialog.askdirectory(title="选择保存目录") try: root.destroy() except Exception: pass if path: return {"path": path} except Exception: pass system = platform.system() # 2) Windows: PowerShell COM Shell.Application if system == "Windows": # 2a) PowerShell COM Shell.Application try: import subprocess ps_cmd = r'$f=(New-Object -ComObject Shell.Application).BrowseForFolder(0,"选择保存目录",0); if($f){$f.Self.Path}' res = subprocess.run( ["powershell", "-NoProfile", "-Command", ps_cmd], capture_output=True, text=True, timeout=60 ) out = (res.stdout or "").strip() if out: return {"path": out} except Exception: pass # 2b) VBScript + cscript(兼容禁用 PowerShell 的环境) try: import subprocess vbs = ( 'Set sh = CreateObject("Shell.Application")\n' 'Set f = sh.BrowseForFolder(0, "选择保存目录", 0)\n' 'If (Not f Is Nothing) Then\n' ' WScript.Echo f.Self.Path\n' 'End If\n' ) with tempfile.NamedTemporaryFile(delete=False, suffix=".vbs") as tf: tf.write(vbs.encode("utf-8")) tf_path = tf.name try: res = subprocess.run( ["cscript", "//nologo", tf_path], capture_output=True, text=True, timeout=60 ) finally: try: os.remove(tf_path) except Exception: pass out = (res.stdout or "").strip() if out: return {"path": out} except Exception: pass # 3) macOS: AppleScript choose folder if system == "Darwin": try: import subprocess script = 'tell application "System Events" to POSIX path of (choose folder with prompt "选择保存目录")' res = subprocess.run( ["osascript", "-e", script], capture_output=True, text=True, timeout=60 ) out = (res.stdout or "").strip() if out: return {"path": out} except Exception: pass # 4) Linux: zenity if system == "Linux": try: import subprocess res = subprocess.run( ["zenity", "--file-selection", "--directory", "--title=选择保存目录"], capture_output=True, text=True, timeout=60 ) out = (res.stdout or "").strip() if out: return {"path": out} except Exception: pass raise HTTPException(status_code=500, detail="无法打开文件夹选择器:所有策略均失败") @app.post("/api/open-dir") def api_open_dir(payload: Dict[str, Any]) -> Dict[str, Any]: """ 打开指定目录;若未传 path,则打开配置中的 output_dir。 Windows 使用 os.startfile,macOS 用 open,Linux 用 xdg-open。 """ path = (payload or {}).get("path") or "" if not path: cfg = load_config() path = cfg.output_dir if not path: raise HTTPException(status_code=400, detail="未提供路径且配置中未设置 output_dir") p = Path(path) try: p.mkdir(parents=True, exist_ok=True) except Exception as e: raise HTTPException(status_code=500, detail=f"创建目录失败: {e}") from e try: if hasattr(os, "startfile"): os.startfile(str(p)) # Windows elif platform.system() == "Darwin": import subprocess subprocess.run(["open", str(p)], check=False) else: import subprocess subprocess.run(["xdg-open", str(p)], check=False) return {"ok": True, "path": str(p)} except Exception as e: raise HTTPException(status_code=500, detail=f"无法打开目录: {e}") from e @app.post("/api/generate/t2i") def api_t2i(req: T2IRequest): cfg = load_config() if not cfg.key: raise HTTPException(status_code=400, detail="尚未配置 key,请先在配置中设置 key。") width = _ensure_x64(req.width or 768) height = _ensure_x64(req.height or 768) try: b64, saved = generate_t2i( cfg, prompt=req.prompt, negative=req.negative or "", width=width, height=height, scale=req.scale, steps=req.steps, sampler=req.sampler, noise_schedule=req.noise_schedule, seed=req.seed, variety=req.variety, decrisp=req.decrisp, cfg_rescale=req.cfg_rescale, ) return {"image_base64": _as_data_uri(b64), "saved_path": saved} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) from e @app.post("/api/generate/i2i") def api_i2i(req: I2IRequest): cfg = load_config() if not cfg.key: raise HTTPException(status_code=400, detail="尚未配置 key,请先在配置中设置 key。") width = _ensure_x64(req.width or 768) height = _ensure_x64(req.height or 768) try: b64, saved = generate_i2i( cfg, positive=req.positive, negative=req.negative or "", image_base64=req.image_base64, width=width, height=height, scale=req.scale, steps=req.steps, sampler=req.sampler, noise_schedule=req.noise_schedule, strength=req.strength or 0.5, noise=req.noise or 0.0, seed=req.seed, variety=req.variety, decrisp=req.decrisp, cfg_rescale=req.cfg_rescale, ) return {"image_base64": _as_data_uri(b64), "saved_path": saved} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) from e @app.post("/api/generate/inpaint") def api_inpaint(req: InpaintRequest): cfg = load_config() if not cfg.key: raise HTTPException(status_code=400, detail="尚未配置 key,请先在配置中设置 key。") width = _ensure_x64(req.width or 768) height = _ensure_x64(req.height or 768) try: b64, saved = generate_inpaint( cfg, positive=req.positive, negative=req.negative or "", image_base64=req.image_base64, mask_base64=req.mask_base64, add_original_image=req.add_original_image, width=width, height=height, scale=req.scale, steps=req.steps, sampler=req.sampler, noise_schedule=req.noise_schedule, strength=req.strength or 0.5, noise=req.noise or 0.0, seed=req.seed, variety=req.variety, decrisp=req.decrisp, cfg_rescale=req.cfg_rescale, ) return {"image_base64": _as_data_uri(b64), "saved_path": saved} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) from e # 前端静态资源(仅保留必要 UI,无教程/仓库链接) _frontend_dir = Path(__file__).resolve().parent.parent / "frontend" # 提示音静态资源映射:/ring -> 项目根/ring 目录(例如 G:\NOVELAI\New NAI\ring) _ring_dir = Path(__file__).resolve().parent.parent / "ring" app.mount("/ring", StaticFiles(directory=_ring_dir, html=False), name="ring") app.mount("/", StaticFiles(directory=_frontend_dir, html=True), name="frontend")