app.py

import os
import sys
import gc
import uuid
import random
import re
import datetime
import json
import tempfile
import locale

# =========================
# 语言检测
# =========================
try:
    system_lang = locale.getdefaultlocale()[0]
    is_chinese = system_lang and system_lang.startswith('zh')
except:
    is_chinese = False

def get_message(key, *args):
    messages = {
        "peft_loaded": ("✅ PEFT 库已加载，LoRA 功能可用。", "✅ PEFT library loaded, LoRA functionality available."),
        "peft_not_detected": ("⚠️ 警告: 未检测到 PEFT 库。LoRA 功能将禁用。", "⚠️ Warning: PEFT library not detected. LoRA functionality will be disabled."),
        "lora_skipped": ("⚠️ [LoRA] 已跳过加载：PEFT 库未安装。", "⚠️ [LoRA] Skipped loading: PEFT library not installed."),
        "transformer_not_loaded": ("⚠️ Transformer 未加载，无法应用 LoRA", "⚠️ Transformer not loaded, cannot apply LoRA"),
        "lora_file_not_exist": ("⚠️ LoRA 文件不存在: {}", "⚠️ LoRA file does not exist: {}"),
        "lora_loading": ("  [LoRA] 正在加载: {} (权重: {} * {} = {:.2f})", "  [LoRA] Loading: {} (weight: {} * {} = {:.2f})"),
        "lora_loaded": ("✅ LoRA 加载成功: {}", "✅ LoRA loaded successfully: {}"),
        "lora_failed": ("❌ LoRA 加载严重失败: {}", "❌ LoRA loading failed critically: {}"),
        "applying_vae": ("正在应用自定义 VAE: {}", "Applying custom VAE: {}"),
        "vae_loaded": ("✅ 自定义 VAE 加载成功", "✅ Custom VAE loaded successfully"),
        "vae_failed": ("⚠️ 自定义 VAE 加载失败: {}", "⚠️ Custom VAE loading failed: {}"),
        "forcing_to_ram": ("  [System] 正在强制将模型搬运至 RAM (请稍候)...", "  [System] Forcing model to RAM (please wait)..."),
        "model_to_ram": ("  [System] 模型已加载至 RAM。", "  [System] Model loaded to RAM."),
        "t2i_low_vram": ("  [T2I] 已启用低显存优化模式", "  [T2I] Low VRAM optimization mode enabled"),
        "t2i_high_end": ("  [T2I] 已启用高端机模式", "  [T2I] High-end GPU mode enabled"),
        "t2i_pipeline_loaded": ("✅ 文生图 Pipeline 加载完成", "✅ Text-to-Image Pipeline loaded"),
        "i2i_pipeline_failed": ("加载图生图 Pipeline 失败：{}", "Failed to load Image-to-Image Pipeline: {}"),
        "i2i_pipeline_loaded": ("✅ 图生图 Pipeline 加载完成", "✅ Image-to-Image Pipeline loaded"),
        "i2i_low_vram": ("  [I2I] 已启用低显存优化模式", "  [I2I] Low VRAM optimization mode enabled"),
        "i2i_high_end": ("  [I2I] 已启用高端机模式", "  [I2I] High-end GPU mode enabled"),
        "generation_stopped": ("🛑 生成已被用户手动停止", "🛑 Generation stopped by user"),
        "upload_image_first": ("⚠️ 请先上传图片！", "⚠️ Please upload an image first!"),
        "i2i_model_failed": ("加载图生图模型失败: {}", "Failed to load Image-to-Image model: {}"),
        "native_inpaint_failed": ("⚠️ 原生 Inpaint 失败 ({})，使用手动混合模式...", "⚠️ Native Inpaint failed ({}), using manual blending mode..."),
        "paint_area": ("⚠️ 请使用画笔在图片上涂抹要修改的区域。", "⚠️ Please use the brush to paint the area to modify on the image."),
        "mask_invalid": ("⚠️ Mask 无效，请确保涂抹了区域。", "⚠️ Mask invalid, please ensure an area is painted."),
        "model_load_failed": ("模型加载失败: {}", "Model loading failed: {}"),
        "inpainting_failed": ("局部重绘失败: {}", "Inpainting failed: {}"),
        "generating": ("生成中", "Generating"),
        "img2img_processing": ("图生图中", "Img2Img processing"),
    }
    zh, en = messages[key]
    return (zh if is_chinese else en).format(*args)

# 环境配置
os.environ.pop("PYTHONHOME", None)
os.environ.pop("PYTHONPATH", None)
os.environ["DIFFUSERS_USE_PEFT_BACKEND"] = "true"
os.environ["PEFT_DEBUG"] = "false"

import torch
import numpy as np
from PIL import Image, ImageFilter, ImageOps, ImageEnhance, ImageDraw

import gradio as gr
from diffusers import (
    ZImagePipeline, 
    ZImageImg2ImgPipeline,
    AutoencoderKL, 
    ZImageTransformer2DModel,
    FlowMatchEulerDiscreteScheduler
)
from transformers import AutoModelForCausalLM, AutoTokenizer
from safetensors.torch import load_file

# =========================
# 检测 PEFT 环境
# =========================
PEFT_AVAILABLE = False
try:
    import peft
    from diffusers.utils import is_peft_available
    if is_peft_available():
        PEFT_AVAILABLE = True
        print(get_message("peft_loaded"))
    else:
        raise ImportError
except ImportError:
    print(get_message("peft_not_detected"))

# =========================
# 双语文本字典
# =========================
TEXT = {
    "zh": {
        "title": "# 🎨 Z-Image-Turbo Low Vram Edition",
        "lang_btn": "EN",
        "tab_generate": "图像生成", "tab_edit": "图片编辑", "tab_img2img": "图生图 (增强版)", "tab_inpaint": "局部重绘",
        "prompt": "Prompt", "prompt_placeholder": "输入你的描述...", "negative_prompt": "负面提示词", "negative_placeholder": "low quality, blurry, bad anatomy",
        "refresh_lora": "🔄 刷新 LoRA", "refresh_model": "🔄 刷新模型", "lora_label": "LoRA", "lora_strength": "LoRA 强度", "lora_weight": "权重",
        "model_section": "### 模型选择/Model Selection", "transformer": "Transformer", "vae": "VAE", "vram_type": "显存类型",
        "vram_low": "24GB以下 (优化模式)", "vram_high": "高端机模式 (>=24GB)", "device": "设备", "num_images": "生成张数",
        "output_format": "输出格式", "width": "宽度", "height": "高度", "steps": "步数", "cfg": "CFG", "seed": "种子", "random_seed": "随机种子",
        "generate": "🚀 生成", "stop": "🛑 停止生成", "gallery": "生成结果", "used_seed": "使用种子",
        "edit_upload": "上传图片", "rotate": "旋转角度 (度)", "crop_x": "裁剪 X (%)", "crop_y": "裁剪 Y (%)", "crop_w": "裁剪宽度 (%)", "crop_h": "裁剪高度 (%)",
        "hflip": "水平翻转", "vflip": "垂直翻转", "edit_btn": "开始编辑", "edited_image": "编辑后的图片",
        "filter": "应用滤镜", "brightness": "亮度调整 (%)", "contrast": "对比度调整 (%)", "saturation": "饱和度调整 (%)",
        "i2i_ref": "上传参考图", "i2i_prompt": "修改提示词", "i2i_ph": "描述你希望图中发生的变化...", "i2i_mode": "Img2Img 模式",
        "i2i_mode_a": "A. 严格保结构（微调风格）", "i2i_mode_b": "B. 强烈听 prompt（允许大改）", "i2i_out_w": "输出宽 (0=自动)", "i2i_out_h": "输出高 (0=自动)",
        "i2i_tip": "**提示：** 宽高都为0时自动保持上传图比例并接近1024。", "i2i_strength": "重绘强度", "i2i_btn": "🎨 开始修改", "i2i_note": "注：使用官方 Z-Image Img2Img 引擎。",
        "inpaint_editor": "绘制 Mask (白色为修改区，黑色为保留区)", "inpaint_tip": "提示：先上传图片，然后用画笔涂抹要修改的区域。", "inpaint_upload": "上传原图并绘制", "inpaint_desc": "📖 使用指南：涂抹区域（白色/彩色）将被重新生成，未涂抹区域保持原样。",
    },
    "en": {
        "title": "# 🎨 Z-Image-Turbo Low Vram Edition", "lang_btn": "中文",
        "tab_generate": "Image Generation", "tab_edit": "Image Editing", "tab_img2img": "Img2Img (Enhanced)", "tab_inpaint": "Inpainting",
        "prompt": "Prompt", "prompt_placeholder": "Enter your description...", "negative_prompt": "Negative Prompt", "negative_placeholder": "low quality, blurry",
        "refresh_lora": "🔄 Refresh LoRA", "refresh_model": "🔄 Refresh Models", "lora_label": "LoRA", "lora_strength": "LoRA Strength", "lora_weight": "Weight",
        "model_section": "### Model Selection", "transformer": "Transformer", "vae": "VAE", "vram_type": "VRAM Type",
        "vram_low": "Under 24GB (Optimized)", "vram_high": "High-End GPU Mode (>=24GB)", "device": "Device", "num_images": "Number of Images",
        "output_format": "Output Format", "width": "Width", "height": "Height", "steps": "Steps", "cfg": "CFG", "seed": "Seed", "random_seed": "Random Seed",
        "generate": "🚀 Generate", "stop": "🛑 Stop Generation", "gallery": "Generated Images", "used_seed": "Used Seed",
        "edit_upload": "Upload Image", "rotate": "Rotation (degrees)", "crop_x": "Crop X (%)", "crop_y": "Crop Y (%)", "crop_w": "Crop Width (%)", "crop_h": "Crop Height (%)",
        "hflip": "Horizontal Flip", "vflip": "Vertical Flip", "edit_btn": "Apply Edit", "edited_image": "Edited Image",
        "filter": "Apply Filter", "brightness": "Brightness (%)", "contrast": "Contrast (%)", "saturation": "Saturation (%)",
        "i2i_ref": "Upload Reference", "i2i_prompt": "Modification Prompt", "i2i_ph": "Describe changes...", "i2i_mode": "Img2Img Mode",
        "i2i_mode_a": "A. Strict Structure (Style tweak)", "i2i_mode_b": "B. Strong Prompt (Allow changes)", "i2i_out_w": "Output Width (0=Auto)", "i2i_out_h": "Output Height (0=Auto)",
        "i2i_tip": "**Tip:** Auto ratio if both 0.", "i2i_strength": "Denoising Strength", "i2i_btn": "🎨 Start Modification", "i2i_note": "Using official Z-Image Img2Img engine.",
        "inpaint_editor": "Draw Mask (White=Modify, Black=Keep)", "inpaint_tip": "Tip: Upload image, then paint area to modify.", "inpaint_upload": "Upload & Paint", "inpaint_desc": "📖 Guide: Painted areas (white/color) will be regenerated. Unpainted areas stay original.",
    }
}

# =========================
# 路径配置
# =========================
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
BASE_SNAPSHOT_DIR = os.path.join(BASE_DIR, "cache", "HF_HOME", "hub", "models--Tongyi-MAI--Z-Image-Turbo", "snapshots", "5f4b9cbb80cc95ba44fe6667dfd75710f7db2947")
if not os.path.exists(BASE_SNAPSHOT_DIR):
    BASE_SNAPSHOT_DIR = os.path.join(BASE_DIR, "ckpts", "Z-Image-Turbo")
    if not os.path.exists(BASE_SNAPSHOT_DIR):
        BASE_SNAPSHOT_DIR = "." 

TRANSFORMER_ROOT = os.path.join(BASE_SNAPSHOT_DIR, "transformer")
TEXT_ENCODER_ROOT = os.path.join(BASE_SNAPSHOT_DIR, "text_encoder")
VAE_ROOT = os.path.join(BASE_SNAPSHOT_DIR, "vae")

MOD_DIR = os.path.join(BASE_DIR, "MOD")
MOD_TRANSFORMER = os.path.join(MOD_DIR, "transformer")
MOD_VAE = os.path.join(MOD_DIR, "vae")
LORA_ROOT = os.path.join(BASE_DIR, "lora")
OUTPUT_DIR = os.path.join(BASE_DIR, "outputs")

for p in [MOD_TRANSFORMER, MOD_VAE, LORA_ROOT, OUTPUT_DIR]:
    os.makedirs(p, exist_ok=True)

pipe_t2i = None
pipe_i2i = None
current_model_config = {"transformer": "default", "vae": "default", "is_low_vram": True}
is_generating_interrupted = False

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32

def auto_flush_vram():
    gc.collect()
    if DEVICE == "cuda":
        torch.cuda.empty_cache()

# =========================
# 核心优化：LoRA 加载逻辑
# =========================
def apply_lora_to_pipeline(pipe_local, lora_choice, lora_alpha, lora_scale=1.0):
    if not PEFT_AVAILABLE:
        print(get_message("lora_skipped"))
        return pipe_local

    if pipe_local is None:
        return pipe_local
    if pipe_local.transformer is None:
        print(get_message("transformer_not_loaded"))
        return pipe_local

    if hasattr(pipe_local, "unload_lora_weights"):
        try:
            pipe_local.unload_lora_weights()
        except Exception:
            pass

    if not lora_choice or lora_choice.lower() == "none":
        return pipe_local

    lora_path = os.path.join(LORA_ROOT, lora_choice)
    if not os.path.exists(lora_path):
        print(get_message("lora_file_not_exist", lora_path))
        return pipe_local

    try:
        raw_alpha = float(lora_alpha)
        effective_alpha = raw_alpha * lora_scale

        if effective_alpha <= 0:
            return pipe_local

        adapter_name = re.sub(r"[^a-zA-Z0-9_]", "_", os.path.splitext(lora_choice)[0])

        print(get_message("lora_loading", lora_choice, raw_alpha, lora_scale, effective_alpha))
        pipe_local.load_lora_weights(
            LORA_ROOT,
            weight_name=lora_choice,
            adapter_name=adapter_name
        )
        pipe_local.set_adapters([adapter_name], adapter_weights=[effective_alpha])
        print(get_message("lora_loaded", adapter_name))

    except Exception as e:
        import traceback
        print(get_message("lora_failed", e))

    return pipe_local

def scan_lora_items():
    if not os.path.isdir(LORA_ROOT):
        return []
    return sorted([f for f in os.listdir(LORA_ROOT) if f.lower().endswith((".safetensors", ".pt", ".pth"))])

def update_prompt_with_lora(prompt, lora_choice, lora_alpha):
    prompt = (prompt or "").strip()
    prompt_clean = re.sub(r"<lora:[^>]+>", "", prompt).strip()
    if lora_choice and lora_choice.lower() != "none":
        try:
            alpha = float(lora_alpha)
        except: alpha = 1.0
        if alpha > 0:
            name = os.path.splitext(lora_choice)[0]
            alpha_str = f"{alpha:.2f}".rstrip("0").rstrip(".")
            return f"{prompt_clean} <lora:{name}:{alpha_str}>"
    return prompt_clean

# =========================
# 模型加载逻辑
# =========================
def load_t2i_pipeline(transformer_choice, vae_choice, is_low_vram):
    global pipe_t2i, current_model_config
    config_key = ("t2i", transformer_choice, vae_choice, is_low_vram)
    if pipe_t2i is not None and current_model_config.get("t2i") == config_key:
        return pipe_t2i

    auto_flush_vram()
    pipe_t2i = None
    
    transformer = ZImageTransformer2DModel.from_pretrained(TRANSFORMER_ROOT, torch_dtype=DTYPE, local_files_only=True)
    if transformer_choice != "default":
        t_path = resolve_model_path(transformer_choice, MOD_TRANSFORMER)
        if t_path:
            if os.path.isdir(t_path):
                custom_t = ZImageTransformer2DModel.from_pretrained(t_path, torch_dtype=DTYPE, local_files_only=True)
                transformer = custom_t
            else:
                state = load_file(t_path, device="cpu")
                processed = {}
                prefix = "model.diffusion_model."
                for k, v in state.items():
                    new_k = k[len(prefix):] if k.startswith(prefix) else k
                    processed[new_k] = v.to(DTYPE)
                transformer.load_state_dict(processed, strict=False)
                del state, processed

    text_encoder = AutoModelForCausalLM.from_pretrained(TEXT_ENCODER_ROOT, torch_dtype=DTYPE, local_files_only=True)
    
    pipe_t2i = ZImagePipeline.from_pretrained(
        BASE_SNAPSHOT_DIR,
        local_files_only=True,
        transformer=transformer,
        text_encoder=text_encoder,
    )
    pipe_t2i.to(dtype=DTYPE)

    if vae_choice != "default":
        v_path = resolve_model_path(vae_choice, MOD_VAE)
        if v_path:
            print(get_message("applying_vae", vae_choice))
            vae_device_map = {"": "cpu"} if is_low_vram else None
            try:
                if os.path.isfile(v_path):
                    with tempfile.TemporaryDirectory() as tmpdir:
                        config_file_path = os.path.join(tmpdir, "config.json")
                        vae_config_dict = dict(pipe_t2i.vae.config)
                        with open(config_file_path, "w", encoding="utf-8") as f:
                            json.dump(vae_config_dict, f, indent=2)
                        try:
                            pipe_t2i.vae = AutoencoderKL.from_single_file(v_path, dtype=DTYPE, config=tmpdir, device_map=vae_device_map)
                        except TypeError:
                            pipe_t2i.vae = AutoencoderKL.from_single_file(v_path, torch_dtype=DTYPE, config=tmpdir, device_map=vae_device_map)
                        print(get_message("vae_loaded"))
                else:
                    pipe_t2i.vae = AutoencoderKL.from_pretrained(v_path, torch_dtype=DTYPE, device_map=vae_device_map)
            except Exception as e:
                print(get_message("vae_failed", e))

    if DEVICE == "cuda":
        if is_low_vram:
            print(get_message("forcing_to_ram"))
            pipe_t2i.to("cpu")
            print(get_message("model_to_ram"))
            pipe_t2i.enable_sequential_cpu_offload()
            print(get_message("t2i_low_vram"))
        else:
            pipe_t2i.to("cuda")
            print(get_message("t2i_high_end"))

    current_model_config["t2i"] = config_key
    print("✅ 文生图 Pipeline 加载完成")
    return pipe_t2i

def load_i2i_pipeline(transformer_choice, vae_choice, is_low_vram):
    global pipe_i2i, current_model_config
    config_key = ("i2i", transformer_choice, vae_choice, is_low_vram)
    if pipe_i2i is not None and current_model_config.get("i2i") == config_key:
        return pipe_i2i

    auto_flush_vram()
    pipe_i2i = None
    
    transformer = ZImageTransformer2DModel.from_pretrained(TRANSFORMER_ROOT, torch_dtype=DTYPE, local_files_only=True)
    if transformer_choice != "default":
        t_path = resolve_model_path(transformer_choice, MOD_TRANSFORMER)
        if t_path:
            if os.path.isdir(t_path):
                custom_t = ZImageTransformer2DModel.from_pretrained(t_path, torch_dtype=DTYPE, local_files_only=True)
                transformer = custom_t
            else:
                state = load_file(t_path, device="cpu")
                processed = {}
                prefix = "model.diffusion_model."
                for k, v in state.items():
                    new_k = k[len(prefix):] if k.startswith(prefix) else k
                    processed[new_k] = v.to(DTYPE)
                transformer.load_state_dict(processed, strict=False)
                del state, processed

    try:
        pipe_i2i = ZImageImg2ImgPipeline.from_pretrained(
            BASE_SNAPSHOT_DIR,
            local_files_only=True,
            transformer=transformer,
        )
    except Exception as e:
        raise gr.Error(f"加载图生图 Pipeline 失败：{str(e)}")
        
    pipe_i2i.to(dtype=DTYPE)

    if vae_choice != "default":
        v_path = resolve_model_path(vae_choice, MOD_VAE)
        if v_path:
            print(get_message("applying_vae", vae_choice))
            vae_device_map = {"": "cpu"} if is_low_vram else None
            try:
                if os.path.isfile(v_path):
                    with tempfile.TemporaryDirectory() as tmpdir:
                        config_file_path = os.path.join(tmpdir, "config.json")
                        vae_config_dict = dict(pipe_i2i.vae.config)
                        with open(config_file_path, "w", encoding="utf-8") as f:
                            json.dump(vae_config_dict, f, indent=2)
                        try:
                            pipe_i2i.vae = AutoencoderKL.from_single_file(v_path, dtype=DTYPE, config=tmpdir, device_map=vae_device_map)
                        except TypeError:
                            pipe_i2i.vae = AutoencoderKL.from_single_file(v_path, torch_dtype=DTYPE, config=tmpdir, device_map=vae_device_map)
                        print(get_message("vae_loaded"))
                else:
                    pipe_i2i.vae = AutoencoderKL.from_pretrained(v_path, torch_dtype=DTYPE, device_map=vae_device_map)
            except Exception as e:
                print(get_message("vae_failed", e))

    if DEVICE == "cuda":
        if is_low_vram:
            print(get_message("forcing_to_ram"))
            pipe_i2i.to("cpu")
            print(get_message("model_to_ram"))
            pipe_i2i.enable_sequential_cpu_offload()
            print(get_message("i2i_low_vram"))
        else:
            pipe_i2i.to("cuda")
            print(get_message("i2i_high_end"))

    current_model_config["i2i"] = config_key
    print("✅ 图生图 Pipeline 加载完成")
    return pipe_i2i

def interrupt_callback(pipe, step, timestep, callback_kwargs):
    global is_generating_interrupted
    if is_generating_interrupted:
        raise gr.Error("🛑 生成已被用户手动停止")
    return callback_kwargs

def scan_model_variants(root_dir):
    if not os.path.isdir(root_dir):
        return []
    items = []
    for name in os.listdir(root_dir):
        path = os.path.join(root_dir, name)
        if os.path.isdir(path):
            if os.path.isfile(os.path.join(path, "config.json")):
                items.append(name)
        elif name.lower().endswith((".safetensors", ".bin")):
            items.append(name)
    return sorted(items)

def get_choices(mod_root):
    return ["default"] + scan_model_variants(mod_root)

def resolve_model_path(choice, mod_root):
    if choice == "default":
        return None
    path = os.path.join(mod_root, choice)
    if os.path.exists(path):
        return path
    return None

def process_mask_for_inpaint(mask_image):
    if mask_image is None:
        return None
    if mask_image.mode == 'RGBA':
        import numpy as np
        mask_array = np.array(mask_image)
        alpha = mask_array[:, :, 3] if mask_array.shape[2] > 3 else None
        rgb = mask_array[:, :, :3]
        rgb_gray = np.dot(rgb, [0.299, 0.587, 0.114])
        if alpha is not None:
            mask_gray = np.where(alpha > 10, 255, 0).astype(np.uint8)
        else:
            mask_gray = np.where(rgb_gray > 10, 255, 0).astype(np.uint8)
        mask = Image.fromarray(mask_gray, mode='L')
    else:
        if mask_image.mode != 'L':
            mask_image = mask_image.convert('L')
        mask = mask_image.point(lambda p: 255 if p > 10 else 0)
    
    if mask.getextrema()[1] == 0: 
        return None
    return mask

# =========================
# 生成与编辑函数
# =========================

def generate_image(prompt, lora_choice, lora_alpha, num_images, image_format,

                   width, height, num_inference_steps, guidance_scale, seed, randomize_seed,

                   transformer_choice, vae_choice, vram_type_str, progress=gr.Progress()):
    global is_generating_interrupted
    is_generating_interrupted = False
    
    is_low_vram = "24GB" in vram_type_str or "Under 24GB" in vram_type_str or "24G以下" in vram_type_str or "24GB以下" in vram_type_str
    
    pipe_local = load_t2i_pipeline(transformer_choice, vae_choice, is_low_vram)
    pipe_local = apply_lora_to_pipeline(pipe_local, lora_choice, lora_alpha)

    if randomize_seed:
        seed = random.randint(0, 2**32 - 1)
    generator = torch.Generator(DEVICE).manual_seed(int(seed))

    date_str = datetime.datetime.now().strftime("%Y-%m-%d")
    day_dir = os.path.join(OUTPUT_DIR, date_str)
    os.makedirs(day_dir, exist_ok=True)

    fmt_map = {"png": ("PNG", "png"), "jpeg": ("JPEG", "jpeg"), "webp": ("WEBP", "webp")}
    pil_fmt, ext = fmt_map[image_format.lower()]

    results = []
    try:
        for _ in progress.tqdm(range(int(num_images)), desc="生成中"):
            if is_generating_interrupted:
                break
            img = pipe_local(
                prompt=prompt.strip(),
                width=width,
                height=height,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                generator=generator,
                callback_on_step_end=interrupt_callback,
            ).images[0]
            filename = os.path.join(day_dir, f"{datetime.datetime.now():%H%M%S}_{uuid.uuid4().hex[:4]}.{ext}")
            img.save(filename, format=pil_fmt)
            results.append(filename)
    finally:
        auto_flush_vram()

    return results, seed

def run_img2img_enhanced(input_image, prompt, negative_prompt, lora_choice, lora_alpha, 

                         num_images, image_format,

                         out_w, out_h, i2i_mode, strength_ui, steps_ui, cfg_ui, 

                         seed, randomize_seed,

                         transformer_choice, vae_choice, vram_type_str, progress=gr.Progress()):
    global is_generating_interrupted
    is_generating_interrupted = False

    is_low_vram = "24GB" in vram_type_str or "Under 24GB" in vram_type_str or "24G以下" in vram_type_str or "24GB以下" in vram_type_str

    if input_image is None:
        raise gr.Error("⚠️ 请先上传图片！")

    try:
        pipe_local = load_i2i_pipeline(transformer_choice, vae_choice, is_low_vram)
    except Exception as e:
        if isinstance(e, gr.Error): raise e
        raise gr.Error(f"加载图生图模型失败: {str(e)}")

    if i2i_mode.startswith("A"):
        lora_scale = 0.35 
        strength = 0.30
        steps = 8
        cfg = 1.0
    else:
        lora_scale = 0.65 
        strength = 0.45
        steps = 6
        cfg = 1.5
    
    pipe_local = apply_lora_to_pipeline(pipe_local, lora_choice, lora_alpha, lora_scale)

    final_strength = strength_ui
    final_steps = int(steps_ui)
    final_cfg = cfg_ui
    
    if randomize_seed:
        seed = random.randint(0, 2**32 - 1)
    generator = torch.Generator(DEVICE).manual_seed(int(seed))

    orig_w, orig_h = input_image.size
    if out_w == 0 or out_h == 0:
        target_size = 1024
        ratio = orig_w / orig_h
        if ratio > 1:
            w, h = target_size, int(target_size / ratio)
        else:
            w, h = int(target_size * ratio), target_size
    else:
        w, h = out_w, out_h
    
    w = (w // 16) * 16
    h = (h // 16) * 16
    input_image = input_image.resize((w, h), Image.LANCZOS)

    date_str = datetime.datetime.now().strftime("%Y-%m-%d")
    day_dir = os.path.join(OUTPUT_DIR, date_str)
    os.makedirs(day_dir, exist_ok=True)

    fmt_map = {"png": ("PNG", "png"), "jpeg": ("JPEG", "jpeg"), "webp": ("WEBP", "webp")}
    pil_fmt, ext = fmt_map[image_format.lower()]

    results = []
    
    try:
        for _ in progress.tqdm(range(int(num_images)), desc="图生图中"):
            if is_generating_interrupted:
                break
            
            img = pipe_local(
                prompt=prompt.strip(),
                negative_prompt=negative_prompt.strip(),
                image=input_image,
                strength=final_strength,
                num_inference_steps=final_steps,
                guidance_scale=final_cfg,
                generator=generator,
                callback_on_step_end=interrupt_callback,
            ).images[0]
            
            filename = os.path.join(day_dir, f"i2i_{datetime.datetime.now():%H%M%S}_{uuid.uuid4().hex[:4]}.{ext}")
            img.save(filename, format=pil_fmt)
            results.append(filename)
    finally:
        auto_flush_vram()

    return results, seed

def run_inpainting(image_editor_data, prompt, negative_prompt, lora_choice, lora_alpha,

                   strength, steps, cfg, seed, randomize_seed,

                   transformer_choice, vae_choice, vram_type_str, progress=gr.Progress()):
    global is_generating_interrupted
    is_generating_interrupted = False

    is_low_vram = "24GB" in vram_type_str or "Under 24GB" in vram_type_str or "24G以下" in vram_type_str or "24GB以下" in vram_type_str

    input_image = None
    mask_layer = None
    
    if isinstance(image_editor_data, dict):
        if 'background' in image_editor_data:
            input_image = image_editor_data['background']
            if image_editor_data.get('layers'):
                mask_layer = image_editor_data['layers'][0]
    elif isinstance(image_editor_data, (tuple, list)):
        input_image = image_editor_data[0]
        mask_layer = image_editor_data[1]
    elif isinstance(image_editor_data, Image.Image):
        input_image = image_editor_data

    if input_image is None:
        raise gr.Error("⚠️ 请先上传图片！")
    
    if input_image.mode == 'RGBA':
        background = Image.new('RGB', input_image.size, (255,255,255))
        background.paste(input_image, (0, 0), input_image)
        input_image = background
    else:
        input_image = input_image.convert("RGB")

    if mask_layer is None:
        raise gr.Error("⚠️ 请使用画笔在图片上涂抹要修改的区域。")
    
    mask = process_mask_for_inpaint(mask_layer)
    if mask is None:
        raise gr.Error("⚠️ Mask 无效，请确保涂抹了区域。")

    try:
        pipe_local = load_i2i_pipeline(transformer_choice, vae_choice, is_low_vram)
    except Exception as e:
        raise gr.Error(f"模型加载失败: {str(e)}")
    
    pipe_local = apply_lora_to_pipeline(pipe_local, lora_choice, lora_alpha, lora_scale=0.6)

    if randomize_seed:
        seed = random.randint(0, 2**32 - 1)
    generator = torch.Generator(DEVICE).manual_seed(int(seed))

    orig_w, orig_h = input_image.size
    if mask.size != (orig_w, orig_h):
        mask = mask.resize((orig_w, orig_h), Image.LANCZOS)

    date_str = datetime.datetime.now().strftime("%Y-%m-%d")
    day_dir = os.path.join(OUTPUT_DIR, date_str)
    os.makedirs(day_dir, exist_ok=True)

    result_img = None

    try:
        try:
            result_img = pipe_local(
                prompt=prompt.strip(),
                negative_prompt=negative_prompt.strip(),
                image=input_image,
                mask_image=mask,
                strength=float(strength),
                num_inference_steps=int(steps),
                guidance_scale=float(cfg),
                generator=generator,
                callback_on_step_end=interrupt_callback
            ).images[0]
        except (TypeError, AttributeError) as e:
            print(f"⚠️ 原生 Inpaint 失败 ({e})，使用手动混合模式...")
            
            img_array = np.array(input_image).astype(np.float32) /255.0
            mask_array = np.array(mask.convert('L')).astype(np.float32) / 255.0
            mask_3d = np.expand_dims(mask_array, axis=2)
            mask_3d = np.repeat(mask_3d,3, axis=2)
            
            noise = np.random.randn(*img_array.shape).astype(np.float32) * 0.1
            inpaint_input_array = img_array * (1 - mask_3d) + (img_array + noise) * mask_3d
            inpaint_input_array = np.clip(inpaint_input_array, 0, 1)
            inpaint_input = Image.fromarray((inpaint_input_array * 255).astype(np.uint8))
            
            generated = pipe_local(
                prompt=prompt.strip(),
                negative_prompt=negative_prompt.strip(),
                image=inpaint_input,
                strength=float(strength),
                num_inference_steps=int(steps),
                guidance_scale=float(cfg),
                generator=generator,
                callback_on_step_end=interrupt_callback
            ).images[0]
            
            if generated.size != (orig_w, orig_h):
                generated = generated.resize((orig_w, orig_h), Image.LANCZOS)
            
            gen_array = np.array(generated).astype(np.float32) / 255.0
            orig_array = np.array(input_image).astype(np.float32) / 255.0
            
            final_array = orig_array * (1 - mask_3d) + gen_array * mask_3d
            final_array = np.clip(final_array, 0, 1)
            result_img = Image.fromarray((final_array * 255).astype(np.uint8))

        filename = os.path.join(day_dir, f"inpaint_{datetime.datetime.now():%H%M%S}_{uuid.uuid4().hex[:4]}.png")
        result_img.save(filename)
        
    except Exception as e:
        if "任务已手动停止" in str(e): raise
        import traceback
        traceback.print_exc()
        raise gr.Error(f"局部重绘失败: {str(e)}")
    finally:
        auto_flush_vram()

    return [result_img], seed

def edit_image(image, angle, x, y, w, h, hflip, vflip, filter_name, brightness, contrast, saturation):
    if image is None:
        return None
    img = image.copy()
    if angle != 0:
        img = img.rotate(angle, expand=True)
    if x or y or w < 100 or h < 100:
        ow, oh = img.size
        left = int(ow * x / 100)
        top = int(oh * y / 100)
        right = int(ow * (x + w) / 100)
        bottom = int(oh * (y + h) / 100)
        img = img.crop((left, top, right, bottom))
    if hflip:
        img = ImageOps.mirror(img)
    if vflip:
        img = ImageOps.flip(img)
    if filter_name:
        filter_map = {
            "模糊": ImageFilter.BLUR, "轮廓": ImageFilter.CONTOUR, "细节": ImageFilter.DETAIL,
            "边缘增强": ImageFilter.EDGE_ENHANCE, "更多边缘增强": ImageFilter.EDGE_ENHANCE_MORE,
            "浮雕": ImageFilter.EMBOSS, "查找边缘": ImageFilter.FIND_EDGES,
            "锐化": ImageFilter.SHARPEN, "平滑": ImageFilter.SMOOTH, "更多平滑": ImageFilter.SMOOTH_MORE,
        }
        f = filter_map.get(filter_name)
        if f:
            img = img.filter(f)
    if brightness != 0:
        img = ImageEnhance.Brightness(img).enhance(1 + brightness / 100)
    if contrast != 0:
        img = ImageEnhance.Contrast(img).enhance(1 + contrast / 100)
    if saturation != 0:
        img = ImageEnhance.Color(img).enhance(1 + saturation / 100)
    return img

# =========================
# Gradio 界面构建
# =========================
TOTAL_VRAM = torch.cuda.get_device_properties(0).total_memory if DEVICE == "cuda" else 0
DEFAULT_PERF_MODE = "高端机模式 (>=24GB)" if TOTAL_VRAM >= 24 * 1024**3 else "24GB以下 (优化模式)"

with gr.Blocks() as demo:
    lang_state = gr.State("zh")

    with gr.Row():
        title_md = gr.Markdown(TEXT["zh"]["title"])
        lang_btn = gr.Button(TEXT["zh"]["lang_btn"], size="sm")

    with gr.Tabs() as tabs:
        with gr.Tab(TEXT["zh"]["tab_generate"]) as tab_gen:
            with gr.Row():
                with gr.Column(scale=4):
                    prompt = gr.Textbox(label=TEXT["zh"]["prompt"], lines=4, placeholder=TEXT["zh"]["prompt_placeholder"])
                    with gr.Row():
                        refresh_lora = gr.Button(TEXT["zh"]["refresh_lora"], size="sm")
                        refresh_model_t2i = gr.Button(TEXT["zh"]["refresh_model"], size="sm")
                    
                    lora_choices = ["None"] + scan_lora_items()
                    lora_drop = gr.Dropdown(label=TEXT["zh"]["lora_label"], choices=lora_choices, value="None")
                    lora_alpha = gr.Slider(0, 2, 1, step=0.05, label=TEXT["zh"]["lora_strength"])

                    model_section_md = gr.Markdown(TEXT["zh"]["model_section"])

                    with gr.Row():
                        transformer_choice = gr.Dropdown(label=TEXT["zh"]["transformer"], choices=get_choices(MOD_TRANSFORMER), value="default")
                        vae_choice = gr.Dropdown(label=TEXT["zh"]["vae"], choices=get_choices(MOD_VAE), value="default")

                    vram_type = gr.Radio(
                        [TEXT["zh"]["vram_low"], TEXT["zh"]["vram_high"]],
                        label=TEXT["zh"]["vram_type"],
                        value=DEFAULT_PERF_MODE
                    )
                    device_ui = gr.Radio(["cuda", "cpu"], label=TEXT["zh"]["device"], value="cuda" if torch.cuda.is_available() else "cpu", visible=False)
                    
                    num_images = gr.Slider(1, 8, 1, step=1, label=TEXT["zh"]["num_images"])
                    image_format = gr.Dropdown(["png", "jpeg", "webp"], value="png", label=TEXT["zh"]["output_format"])

                    with gr.Row():
                        width = gr.Slider(512, 2048, 1024, step=64, label=TEXT["zh"]["width"])
                        height = gr.Slider(512, 2048, 1024, step=64, label=TEXT["zh"]["height"])
                    num_inference_steps = gr.Slider(1, 50, 10, step=1, label=TEXT["zh"]["steps"])
                    guidance_scale = gr.Slider(0, 10, 0, step=0.1, label=TEXT["zh"]["cfg"])
                    seed = gr.Number(label=TEXT["zh"]["seed"], value=42, precision=0)
                    randomize_seed = gr.Checkbox(label=TEXT["zh"]["random_seed"], value=True)

                    with gr.Row():
                        generate_btn = gr.Button(TEXT["zh"]["generate"], variant="primary", size="lg")
                        stop_btn = gr.Button(TEXT["zh"]["stop"], variant="stop", size="lg", interactive=False)

                with gr.Column(scale=6):
                    gallery = gr.Gallery(label=TEXT["zh"]["gallery"], columns=2, height="80vh")
                    used_seed = gr.Number(label=TEXT["zh"]["used_seed"], interactive=False)

        with gr.Tab(TEXT["zh"]["tab_edit"]) as tab_edit:
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(label=TEXT["zh"]["edit_upload"], type="pil")
                    with gr.Group():
                        rotate_angle = gr.Slider(-360, 360, 0, step=1, label=TEXT["zh"]["rotate"])
                        crop_x = gr.Slider(0, 100, 0, step=1, label=TEXT["zh"]["crop_x"])
                        crop_y = gr.Slider(0, 100, 0, step=1, label=TEXT["zh"]["crop_y"])
                        crop_width = gr.Slider(0, 100, 100, step=1, label=TEXT["zh"]["crop_w"])
                        crop_height = gr.Slider(0, 100, 100, step=1, label=TEXT["zh"]["crop_h"])
                        flip_horizontal = gr.Checkbox(label=TEXT["zh"]["hflip"])
                        flip_vertical = gr.Checkbox(label=TEXT["zh"]["vflip"])
                    edit_btn = gr.Button(TEXT["zh"]["edit_btn"], variant="primary")

                with gr.Column():
                    edited_image_output = gr.Image(label=TEXT["zh"]["edited_image"], type="pil")
                    with gr.Group():
                        apply_filter = gr.Dropdown(
                            ["模糊", "轮廓", "细节", "边缘增强", "更多边缘增强", "浮雕", "查找边缘", "锐化", "平滑", "更多平滑"],
                            label=TEXT["zh"]["filter"]
                        )
                        brightness = gr.Slider(-100, 100, 0, step=1, label=TEXT["zh"]["brightness"])
                        contrast = gr.Slider(-100, 100, 0, step=1, label=TEXT["zh"]["contrast"])
                        saturation = gr.Slider(-100, 100, 0, step=1, label=TEXT["zh"]["saturation"])

        with gr.Tab(TEXT["zh"]["tab_img2img"]) as tab_img2img:
            i2i_status_md = gr.Markdown(TEXT["zh"]["i2i_note"])
            with gr.Row():
                with gr.Column(scale=4):
                    i2i_image_input = gr.Image(label=TEXT["zh"]["i2i_ref"], type="pil")
                    
                    i2i_prompt = gr.Textbox(label=TEXT["zh"]["i2i_prompt"], lines=3, placeholder=TEXT["zh"]["i2i_ph"])
                    i2i_negative_prompt = gr.Textbox(label=TEXT["zh"]["negative_prompt"], lines=2, placeholder=TEXT["zh"]["negative_placeholder"])
                    
                    with gr.Row():
                        i2i_refresh_lora = gr.Button(TEXT["zh"]["refresh_lora"], size="sm")
                        i2i_refresh_model = gr.Button(TEXT["zh"]["refresh_model"], size="sm")
                    
                    i2i_lora_choices = ["None"] + scan_lora_items()
                    i2i_lora_drop = gr.Dropdown(label=TEXT["zh"]["lora_label"], choices=i2i_lora_choices, value="None")
                    i2i_lora_alpha = gr.Slider(0, 2, 1, step=0.05, label=TEXT["zh"]["lora_strength"])

                    with gr.Accordion(TEXT["zh"]["model_section"], open=False):
                        i2i_transformer_choice = gr.Dropdown(label=TEXT["zh"]["transformer"], choices=get_choices(MOD_TRANSFORMER), value="default")
                        i2i_vae_choice = gr.Dropdown(label=TEXT["zh"]["vae"], choices=get_choices(MOD_VAE), value="default")
                        i2i_vram_type = gr.Radio(
                            [TEXT["zh"]["vram_low"], TEXT["zh"]["vram_high"]],
                            label=TEXT["zh"]["vram_type"],
                            value=DEFAULT_PERF_MODE
                        )

                    i2i_mode = gr.Radio(
                        [TEXT["zh"]["i2i_mode_a"], TEXT["zh"]["i2i_mode_b"]],
                        label=TEXT["zh"]["i2i_mode"],
                        value=TEXT["zh"]["i2i_mode_a"]
                    )
                    
                    with gr.Row():
                        i2i_out_w = gr.Slider(0, 2048, 0, step=16, label=TEXT["zh"]["i2i_out_w"])
                        i2i_out_h = gr.Slider(0, 2048, 0, step=16, label=TEXT["zh"]["i2i_out_h"])
                    i2i_tip_md = gr.Markdown(TEXT["zh"]["i2i_tip"])
                    
                    i2i_strength = gr.Slider(0.1, 1.0, 0.4, step=0.05, label=TEXT["zh"]["i2i_strength"])
                    i2i_steps = gr.Slider(1, 50, 6, step=1, label=TEXT["zh"]["steps"])
                    i2i_cfg = gr.Slider(0.0, 5.0, 1.0, step=0.1, label=TEXT["zh"]["cfg"])
                    
                    i2i_num_images = gr.Slider(1, 4, 1, step=1, label=TEXT["zh"]["num_images"])
                    i2i_image_format = gr.Dropdown(["png", "jpeg", "webp"], value="png", label=TEXT["zh"]["output_format"])
                    i2i_seed = gr.Number(label=TEXT["zh"]["seed"], value=42, precision=0)
                    i2i_randomize_seed = gr.Checkbox(label=TEXT["zh"]["random_seed"], value=True)

                    with gr.Row():
                        i2i_generate_btn = gr.Button(TEXT["zh"]["i2i_btn"], variant="primary", size="lg")
                        i2i_stop_btn = gr.Button(TEXT["zh"]["stop"], variant="stop", size="lg", interactive=False)

                with gr.Column(scale=6):
                    i2i_gallery = gr.Gallery(label=TEXT["zh"]["gallery"], columns=2, height="80vh")
                    i2i_used_seed = gr.Number(label=TEXT["zh"]["used_seed"], interactive=False)

        with gr.Tab(TEXT["zh"]["tab_inpaint"]) as tab_inpaint:
            with gr.Row():
                with gr.Column(scale=4):
                    inpaint_editor = gr.ImageEditor(
                        label=TEXT["zh"]["inpaint_upload"],
                        type="pil",
                        layers=True,
                        eraser=True,
                        brush=gr.Brush(colors=["#FFFFFF", "#000000", "#FF0000"], color_mode="fixed")
                    )
                    inpaint_tip_md = gr.Markdown(TEXT["zh"]["inpaint_desc"])
                    
                    inpaint_prompt = gr.Textbox(label=TEXT["zh"]["i2i_prompt"], lines=3, placeholder=TEXT["zh"]["i2i_ph"])
                    inpaint_negative_prompt = gr.Textbox(label=TEXT["zh"]["negative_prompt"], lines=2, placeholder=TEXT["zh"]["negative_placeholder"])

                    with gr.Row():
                        inpaint_refresh_lora = gr.Button(TEXT["zh"]["refresh_lora"], size="sm")
                        inpaint_refresh_model = gr.Button(TEXT["zh"]["refresh_model"], size="sm")
                    
                    inpaint_lora_choices = ["None"] + scan_lora_items()
                    inpaint_lora_drop = gr.Dropdown(label=TEXT["zh"]["lora_label"], choices=inpaint_lora_choices, value="None")
                    inpaint_lora_alpha = gr.Slider(0, 2, 1, step=0.05, label=TEXT["zh"]["lora_strength"])

                    with gr.Accordion(TEXT["zh"]["model_section"], open=False):
                        inpaint_transformer_choice = gr.Dropdown(label=TEXT["zh"]["transformer"], choices=get_choices(MOD_TRANSFORMER), value="default")
                        inpaint_vae_choice = gr.Dropdown(label=TEXT["zh"]["vae"], choices=get_choices(MOD_VAE), value="default")
                        inpaint_vram_type = gr.Radio(
                            [TEXT["zh"]["vram_low"], TEXT["zh"]["vram_high"]],
                            label=TEXT["zh"]["vram_type"],
                            value=DEFAULT_PERF_MODE
                        )
                    
                    inpaint_strength = gr.Slider(0.1, 1.0, 0.7, step=0.05, label=TEXT["zh"]["i2i_strength"])
                    inpaint_steps = gr.Slider(1, 50, 8, step=1, label=TEXT["zh"]["steps"])
                    inpaint_cfg = gr.Slider(0.0, 5.0, 1.0, step=0.1, label=TEXT["zh"]["cfg"])
                    
                    inpaint_seed = gr.Number(label=TEXT["zh"]["seed"], value=42, precision=0)
                    inpaint_randomize_seed = gr.Checkbox(label=TEXT["zh"]["random_seed"], value=True)

                    with gr.Row():
                        inpaint_generate_btn = gr.Button(TEXT["zh"]["i2i_btn"], variant="primary", size="lg")
                        inpaint_stop_btn = gr.Button(TEXT["zh"]["stop"], variant="stop", size="lg", interactive=False)

                with gr.Column(scale=6):
                    inpaint_gallery = gr.Gallery(label=TEXT["zh"]["gallery"], columns=2, height="80vh")
                    inpaint_used_seed = gr.Number(label=TEXT["zh"]["used_seed"], interactive=False)

    def switch_language_full(lang):
        new_lang = "en" if lang == "zh" else "zh"
        t = TEXT[new_lang]
        
        # 修复：根据硬件显存大小，决定当前应该选哪个语言版本的选项
        is_low_vram_hardware = TOTAL_VRAM < 24 * 1024**3
        current_vram_val = t['vram_low'] if is_low_vram_hardware else t['vram_high']
        
        # 修复：更新显存选项的值，不仅仅是选项列表
        return (
            new_lang, t['title'], t['lang_btn'],
            gr.update(label=t['tab_generate']), gr.update(label=t['tab_edit']), 
            gr.update(label=t['tab_img2img']), gr.update(label=t['tab_inpaint']),
            gr.update(label=t['prompt'], placeholder=t['prompt_placeholder']),
            gr.update(value=t['refresh_lora']), gr.update(value=t['refresh_model']),
            gr.update(label=t['lora_label']), gr.update(label=t['lora_strength']),
            t['model_section'], gr.update(label=t['transformer']), gr.update(label=t['vae']),
            # T2I VRAM: 更新选项和值
            gr.update(label=t['vram_type'], choices=[t['vram_low'], t['vram_high']], value=current_vram_val), 
            gr.update(label=t['device']),
            gr.update(label=t['num_images']), gr.update(label=t['output_format']),
            gr.update(label=t['width']), gr.update(label=t['height']),
            gr.update(label=t['steps']), gr.update(label=t['cfg']),
            gr.update(label=t['seed']), gr.update(label=t['random_seed']),
            gr.update(value=t['generate']), gr.update(value=t['stop']),
            gr.update(label=t['gallery']), gr.update(label=t['used_seed']),
            
            gr.update(label=t['edit_upload']),
            gr.update(label=t['rotate']), gr.update(label=t['crop_x']), gr.update(label=t['crop_y']),
            gr.update(label=t['crop_w']), gr.update(label=t['crop_h']),
            gr.update(label=t['hflip']), gr.update(label=t['vflip']),
            gr.update(value=t['edit_btn']), gr.update(label=t['edited_image']),
            gr.update(label=t['filter']), gr.update(label=t['brightness']), gr.update(label=t['contrast']), gr.update(label=t['saturation']),
            
            gr.update(value=t['i2i_note']),
            gr.update(label=t['i2i_ref']),
            gr.update(label=t['i2i_prompt'], placeholder=t['i2i_ph']),
            gr.update(label=t['negative_prompt'], placeholder=t['negative_placeholder']),
            gr.update(value=t['refresh_lora']), gr.update(value=t['refresh_model']),
            gr.update(label=t['lora_label']), gr.update(label=t['lora_strength']),
            gr.update(label=t['transformer']), gr.update(label=t['vae']),
            # Img2Img VRAM: 更新选项和值
            gr.update(label=t['vram_type'], choices=[t['vram_low'], t['vram_high']], value=current_vram_val),
            gr.update(label=t['i2i_mode'], choices=[t['i2i_mode_a'], t['i2i_mode_b']]),
            gr.update(label=t['i2i_out_w']), gr.update(label=t['i2i_out_h']),
            gr.update(value=t['i2i_tip']),
            gr.update(label=t['i2i_strength']),
            gr.update(label=t['steps']), gr.update(label=t['cfg']),
            gr.update(label=t['num_images']), gr.update(label=t['output_format']),
            gr.update(label=t['seed']), gr.update(label=t['random_seed']),
            gr.update(value=t['i2i_btn']), gr.update(value=t['stop']),
            gr.update(label=t['gallery']), gr.update(label=t['used_seed']),

            gr.update(label=t['inpaint_upload']),
            gr.update(value=t['inpaint_desc']),
            gr.update(label=t['i2i_prompt'], placeholder=t['i2i_ph']),
            gr.update(label=t['negative_prompt'], placeholder=t['negative_placeholder']),
            gr.update(value=t['refresh_lora']), gr.update(value=t['refresh_model']),
            gr.update(label=t['lora_label']), gr.update(label=t['lora_strength']),
            gr.update(label=t['transformer']), gr.update(label=t['vae']),
            # Inpaint VRAM: 更新选项和值
            gr.update(label=t['vram_type'], choices=[t['vram_low'], t['vram_high']], value=current_vram_val),
            gr.update(label=t['i2i_strength']),
            gr.update(label=t['steps']), gr.update(label=t['cfg']),
            gr.update(label=t['seed']), gr.update(label=t['random_seed']),
            gr.update(value=t['i2i_btn']), gr.update(value=t['stop']),
            gr.update(label=t['gallery']), gr.update(label=t['used_seed']),
        )

    lang_btn.click(
        fn=switch_language_full,
        inputs=lang_state,
        outputs=[
            lang_state, title_md, lang_btn,
            tab_gen, tab_edit, tab_img2img, tab_inpaint,
            prompt, refresh_lora, refresh_model_t2i, lora_drop, lora_alpha, model_section_md,
            transformer_choice, vae_choice, vram_type, device_ui, num_images, image_format,
            width, height, num_inference_steps, guidance_scale, seed, randomize_seed,
            generate_btn, stop_btn, gallery, used_seed,
            image_input, rotate_angle, crop_x, crop_y, crop_width, crop_height,
            flip_horizontal, flip_vertical, edit_btn, edited_image_output,
            apply_filter, brightness, contrast, saturation,
            i2i_status_md, i2i_image_input, i2i_prompt, i2i_negative_prompt, 
            i2i_refresh_lora, i2i_refresh_model, i2i_lora_drop, i2i_lora_alpha,
            i2i_transformer_choice, i2i_vae_choice, i2i_vram_type, i2i_mode,
            i2i_out_w, i2i_out_h, i2i_tip_md,
            i2i_strength, i2i_steps, i2i_cfg,
            i2i_num_images, i2i_image_format, i2i_seed, i2i_randomize_seed,
            i2i_generate_btn, i2i_stop_btn, i2i_gallery, i2i_used_seed,
            inpaint_editor, inpaint_tip_md,
            inpaint_prompt, inpaint_negative_prompt,
            inpaint_refresh_lora, inpaint_refresh_model, inpaint_lora_drop, inpaint_lora_alpha,
            inpaint_transformer_choice, inpaint_vae_choice, inpaint_vram_type,
            inpaint_strength, inpaint_steps, inpaint_cfg,
            inpaint_seed, inpaint_randomize_seed,
            inpaint_generate_btn, inpaint_stop_btn, inpaint_gallery, inpaint_used_seed
        ]
    )

    refresh_lora.click(fn=scan_lora_items, outputs=[lora_drop, i2i_lora_drop, inpaint_lora_drop])
    lora_drop.change(update_prompt_with_lora, [prompt, lora_drop, lora_alpha], prompt)

    def refresh_models_t2i():
        return gr.update(choices=get_choices(MOD_TRANSFORMER)), gr.update(choices=get_choices(MOD_VAE))
    refresh_model_t2i.click(fn=refresh_models_t2i, outputs=[transformer_choice, vae_choice])

    def start_gen(): return gr.update(interactive=False), gr.update(interactive=True)
    def end_gen(): return gr.update(interactive=True), gr.update(interactive=False)
    def trigger_stop():
        global is_generating_interrupted
        is_generating_interrupted = True

    generate_event = generate_btn.click(fn=start_gen, outputs=[generate_btn, stop_btn]).then(
        fn=generate_image,
        inputs=[prompt, lora_drop, lora_alpha, num_images, image_format,
                width, height, num_inference_steps, guidance_scale, seed, randomize_seed,
                transformer_choice, vae_choice, vram_type], 
        outputs=[gallery, used_seed]
    ).then(fn=end_gen, outputs=[generate_btn, stop_btn])

    stop_btn.click(fn=trigger_stop).then(fn=end_gen, outputs=[generate_btn, stop_btn], cancels=[generate_event])

    i2i_refresh_lora.click(fn=scan_lora_items, outputs=[lora_drop, i2i_lora_drop, inpaint_lora_drop])
    i2i_lora_drop.change(update_prompt_with_lora, [i2i_prompt, i2i_lora_drop, i2i_lora_alpha], i2i_prompt)

    def refresh_models_i2i():
        return gr.update(choices=get_choices(MOD_TRANSFORMER)), gr.update(choices=get_choices(MOD_VAE))
    i2i_refresh_model.click(fn=refresh_models_i2i, outputs=[i2i_transformer_choice, i2i_vae_choice])

    def start_i2i(): return gr.update(interactive=False), gr.update(interactive=True)
    def end_i2i(): return gr.update(interactive=True), gr.update(interactive=False)

    i2i_generate_event = i2i_generate_btn.click(fn=start_i2i, outputs=[i2i_generate_btn, i2i_stop_btn]).then(
        fn=run_img2img_enhanced,
        inputs=[i2i_image_input, i2i_prompt, i2i_negative_prompt, i2i_lora_drop, i2i_lora_alpha, 
                i2i_num_images, i2i_image_format,
                i2i_out_w, i2i_out_h, i2i_mode, i2i_strength, i2i_steps, i2i_cfg,
                i2i_seed, i2i_randomize_seed,
                i2i_transformer_choice, i2i_vae_choice, i2i_vram_type], 
        outputs=[i2i_gallery, i2i_used_seed]
    ).then(fn=end_i2i, outputs=[i2i_generate_btn, i2i_stop_btn])

    i2i_stop_btn.click(fn=trigger_stop).then(fn=end_i2i, outputs=[i2i_generate_btn, i2i_stop_btn], cancels=[i2i_generate_event])

    inpaint_refresh_lora.click(fn=scan_lora_items, outputs=[lora_drop, i2i_lora_drop, inpaint_lora_drop])
    inpaint_lora_drop.change(update_prompt_with_lora, [inpaint_prompt, inpaint_lora_drop, inpaint_lora_alpha], inpaint_prompt)

    def refresh_models_inpaint():
        return gr.update(choices=get_choices(MOD_TRANSFORMER)), gr.update(choices=get_choices(MOD_VAE))
    inpaint_refresh_model.click(fn=refresh_models_inpaint, outputs=[inpaint_transformer_choice, inpaint_vae_choice])

    def start_inpaint(): return gr.update(interactive=False), gr.update(interactive=True)
    def end_inpaint(): return gr.update(interactive=True), gr.update(interactive=False)

    inpaint_generate_event = inpaint_generate_btn.click(fn=start_inpaint, outputs=[inpaint_generate_btn, inpaint_stop_btn]).then(
        fn=run_inpainting,
        inputs=[inpaint_editor, inpaint_prompt, inpaint_negative_prompt, inpaint_lora_drop, inpaint_lora_alpha,
                inpaint_strength, inpaint_steps, inpaint_cfg,
                inpaint_seed, inpaint_randomize_seed,
                inpaint_transformer_choice, inpaint_vae_choice, inpaint_vram_type], 
        outputs=[inpaint_gallery, inpaint_used_seed]
    ).then(fn=end_inpaint, outputs=[inpaint_generate_btn, inpaint_stop_btn])

    inpaint_stop_btn.click(fn=trigger_stop).then(fn=end_inpaint, outputs=[inpaint_generate_btn, inpaint_stop_btn], cancels=[inpaint_generate_event])

    edit_btn.click(
        fn=edit_image,
        inputs=[image_input, rotate_angle, crop_x, crop_y, crop_width, crop_height,
                flip_horizontal, flip_vertical, apply_filter, brightness, contrast, saturation],
        outputs=edited_image_output
    )

if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch(show_error=True)