glut.py

import os
import gc
import time
from PIL import Image
import math
from mmgp import offload
import torch
import numpy as np
import gradio as gr
import socket
import psutil
import random
import argparse
import datetime
from diffusers import ZImagePipeline, ZImageTransformer2DModel, ZImageImg2ImgPipeline
from diffusers.utils import load_image
from videox_fun.utils.utils import get_image_latent
from videox_fun.models import ZImageControlTransformer2DModel
from videox_fun.pipeline import ZImageControlPipeline
from transformers import Qwen3Model
from safetensors.torch import load_file


parser = argparse.ArgumentParser() 
parser.add_argument("--server_name", type=str, default="127.0.0.1", help="IP地址，局域网访问改为0.0.0.0")
parser.add_argument("--server_port", type=int, default=7891, help="使用端口")
parser.add_argument("--share", action="store_true", help="是否启用gradio共享")
parser.add_argument("--mcp_server", action="store_true", help="是否启用mcp服务")
parser.add_argument("--compile", action="store_true", help="是否启用compile加速")
parser.add_argument("--res_vram", type=int, default=2000, help="显存保留，单位MB。数值越大，占用显存越小，速度越慢")
args = parser.parse_args()

print(" 启动中，请耐心等待 bilibili@十字鱼 https://space.bilibili.com/893892")
print(f'\033[32mPytorch版本：{torch.__version__}\033[0m')
if torch.cuda.is_available():
    device = "cuda" 
    print(f'\033[32m显卡型号：{torch.cuda.get_device_name()}\033[0m')
    total_vram_in_gb = torch.cuda.get_device_properties(0).total_memory / 1073741824
    print(f'\033[32m显存大小：{total_vram_in_gb:.2f}GB\033[0m')
    mem = psutil.virtual_memory()
    print(f'\033[32m内存大小：{mem.total/1073741824:.2f}GB\033[0m')
    if torch.cuda.get_device_capability()[0] >= 8:
        print(f'\033[32m支持BF16\033[0m')
        dtype = torch.bfloat16
    else:
        print(f'\033[32m不支持BF16，尝试使用FP32\033[0m')
        dtype = torch.float32
else:
    print(f'\033[32mCUDA不可用，请检查\033[0m')
    device = "cpu"

# 启用 CUDA 加速优化
if torch.cuda.is_available():
    torch.backends.cudnn.benchmark = True  # 自动寻找最优卷积算法
    torch.backends.cuda.matmul.allow_tf32 = True  # 允许 TF32 矩阵乘法
    torch.backends.cudnn.allow_tf32 = True  # 允许 TF32 加速

os.makedirs("outputs", exist_ok=True)
repo_id = "./models/Z-Image-Turbo"
if torch.cuda.is_available():
    budgets = int(torch.cuda.get_device_properties(0).total_memory/1048576 - args.res_vram)
else:
    budgets = 0
stop_generation = False
mode_loaded = None
pipe = None
mmgp = None
lora_loaded = None
lora_loaded_weights = None
# 提示词缓存
prompt_embeds_cache = {
    "key": None,
    "prompt_embeds": None,
    "prompt_embeds_mask": None,
}
lora_dir = "models/lora"
if os.path.exists(lora_dir):
    lora_files = [f for f in os.listdir(lora_dir) if f.endswith(".safetensors")]
    lora_choices = sorted(lora_files)
else:
    lora_choices = []

def load_model(mode, lora_dropdown, lora_weights):
    global pipe, mmgp
    text_encoder = offload.fast_load_transformers_model(
        f"{repo_id}/text_encoder/mmgp.safetensors",
        do_quantize=False,
        modelClass=Qwen3Model,
        forcedConfigPath=f"{repo_id}/text_encoder/config.json",
    )
    #text_encoder._dtype = dtype 
    if mode == "t2i":
        if pipe is not None:
            mmgp.release()
        transformer = offload.fast_load_transformers_model(
            f"{repo_id}/transformer/mmgp-image-turbo.safetensors",
            do_quantize=False,
            modelClass=ZImageTransformer2DModel,
            forcedConfigPath=f"{repo_id}/transformer/config.json",
        )
        pipe = ZImagePipeline.from_pretrained(
            repo_id, 
            text_encoder=text_encoder,
            transformer=transformer,
            torch_dtype=dtype,
            low_cpu_mem_usage=False, 
        )
        load_lora(lora_dropdown, lora_weights)
    elif mode == "t2i_image":
        if pipe is not None:
            mmgp.release()
        transformer = offload.fast_load_transformers_model(
            f"{repo_id}/transformer/mmgp-image.safetensors",
            do_quantize=False,
            modelClass=ZImageTransformer2DModel,
            forcedConfigPath=f"{repo_id}/transformer/config.json",
        )
        pipe = ZImagePipeline.from_pretrained(
            repo_id, 
            text_encoder=text_encoder,
            transformer=transformer,
            torch_dtype=dtype,
            low_cpu_mem_usage=False, 
        )
        load_lora(lora_dropdown, lora_weights)
    elif mode == "con":
        if pipe is not None:
            mmgp.release()
        """transformer = ZImageControlTransformer2DModel.from_pretrained(
            repo_id, 
            subfolder="transformer",
            torch_dtype=dtype,
        )
        state_dict = load_file("./models/Z-Image-Turbo-Fun-Controlnet-Union-2.0.safetensors")
        state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
        m, u = transformer.load_state_dict(state_dict, strict=False)
        print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")"""
        transformer = offload.fast_load_transformers_model(
            f"{repo_id}/transformer/mmgp-con.safetensors",
            do_quantize=False,
            modelClass=ZImageControlTransformer2DModel,
            forcedConfigPath=f"{repo_id}/transformer/config.json",
        )
        pipe = ZImageControlPipeline.from_pretrained(
            repo_id, 
            text_encoder=text_encoder,
            transformer=transformer,
            torch_dtype=dtype,
            low_cpu_mem_usage=False,
        )
        load_lora(lora_dropdown, lora_weights)
    mmgp = offload.all(
        pipe, 
        pinnedMemory= ["text_encoder", "transformer"],
        budgets={'*': budgets}, 
        extraModelsToQuantize = ["text_encoder"],
        compile=True if args.compile else False,
    )
    """offload.save_model(
        model=pipe.transformer, 
        file_path=f"{repo_id}/transformer/mmgp-image.safetensors", 
        config_file_path=f"{repo_id}/transformer/config.json",
    )"""
    """offload.save_model(
        model=pipe.text_encoder, 
        file_path=f"{repo_id}/text_encoder/mmgp.safetensors", 
        config_file_path=f"{repo_id}/text_encoder/config.json",
    )"""


def load_lora(lora_dropdown, lora_weights):
    if lora_dropdown != []:
        global pipe
        adapter_names = []
        weightss = []
        weights = [float(w) for w in lora_weights.split(',')] if lora_weights else []
        for idx, lora_name in enumerate(lora_dropdown):
                adapter_name = os.path.splitext(os.path.basename(lora_name))[0]
                adapter_names.append(adapter_name)
                weight = weights[idx] if idx < len(weights) else 1.0
                weightss.append(weight)
                pipe.load_lora_weights(f"models/lora/{lora_name}", adapter_name=adapter_name)
                print(f"✅ 已加载LoRA模型: {lora_name} (权重: {weight})")
        pipe.set_adapters(adapter_names, adapter_weights=weightss)
        print("LoRA加载完成")

# 解决冲突端口（感谢licyk酱提供的代码~）
def find_port(port: int) -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.settimeout(1)
        if s.connect_ex(("localhost", port)) == 0:
            print(f"端口 {port} 已被占用，正在寻找可用端口...")
            return find_port(port=port + 1)
        else:
            return port
        

def exchange_width_height(width, height):
    return height, width, "✅ 宽高交换完毕"


def adjust_width_height(image):
    image_width, image_height = image.size
    vae_width, vae_height = calculate_dimensions(1024*1024, image_width / image_height)
    calculated_height = vae_height // 32 * 32
    calculated_width = vae_width // 32 * 32
    return int(calculated_width), int(calculated_height), "✅ 根据图片调整宽高"


def calculate_dimensions(target_area, ratio):
    width = math.sqrt(target_area * ratio)
    height = width / ratio
    width = round(width / 32) * 32
    height = round(height / 32) * 32
    return width, height


def stop_generate():
    global stop_generation
    stop_generation = True
    return "🛑 等待生成中止"


def scale_resolution_1_5(width, height):
    """
    将宽度和高度都放大1.5倍，并按照16的倍数向下取整
    """
    new_width = int(width * 1.5) // 16 * 16
    new_height = int(height * 1.5) // 16 * 16
    return new_width, new_height, "✅ 分辨率已调整为1.5倍"


def get_cached_prompt_embeds(prompt, negative_prompt=None):
    """获取缓存的 prompt_embeds，避免重复编码"""
    global prompt_embeds_cache, prompt_cache
    
    # 生成缓存键（包含 negative_prompt）
    cache_key = (prompt, negative_prompt)
    
    # 检查缓存
    if prompt_embeds_cache["key"] == cache_key and prompt_embeds_cache["prompt_embeds"] is not None:
        print("📦 使用缓存的 prompt_embeds")
        return prompt_embeds_cache["prompt_embeds"], prompt_embeds_cache["prompt_embeds_mask"]
    
    # 编码新的提示词
    print("🔄 编码提示词...")
    with torch.inference_mode():
        if negative_prompt:
            prompt_embeds, prompt_embeds_mask = pipe.encode_prompt(prompt, negative_prompt=negative_prompt)
        else:
            prompt_embeds, prompt_embeds_mask = pipe.encode_prompt(prompt)
    
    # 更新缓存
    prompt_embeds_cache["key"] = cache_key
    prompt_embeds_cache["prompt_embeds"] = prompt_embeds
    prompt_embeds_cache["prompt_embeds_mask"] = prompt_embeds_mask
    
    return prompt_embeds, prompt_embeds_mask


def generate_t2i(
    prompt, 
    width, 
    height, 
    num_inference_steps, 
    batch_images, 
    seed_param, 
    lora_dropdown, 
    lora_weights,
):
    global stop_generation, mode_loaded, lora_loaded, lora_loaded_weights
    if mode_loaded != "t2i" or lora_loaded != lora_dropdown or lora_loaded_weights != lora_weights:
        load_model("t2i", lora_dropdown, lora_weights)
        mode_loaded = "t2i"
        lora_loaded, lora_loaded_weights = lora_dropdown, lora_weights
        # 模型切换时清除缓存
        global prompt_embeds_cache
        prompt_embeds_cache = {"key": None, "prompt_embeds": None, "prompt_embeds_mask": None}
    results = []
    if seed_param < 0:
        seed = random.randint(0, np.iinfo(np.int32).max)
    else:
        seed = seed_param
    prompt_embeds, _ = get_cached_prompt_embeds(prompt)
    total_start_time = time.time()
    inference_times = []
    for i in range(batch_images):
        if stop_generation:
            stop_generation = False
            yield results, f"✅ 生成已中止，最后种子数{seed+i-1}"
            break
        img_start_time = time.time()
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"outputs/{timestamp}.png"
        with torch.inference_mode():
            output = pipe(
                height=height,
                width=width,
                num_inference_steps=num_inference_steps, 
                guidance_scale=0.0, 
                generator=torch.Generator().manual_seed(seed+i),
                prompt_embeds=prompt_embeds,
            )
        image = output.images[0]
        image.save(filename)
        results.append(image)
        
        # 计算单张图生成时间
        img_time = time.time() - img_start_time
        inference_times.append(img_time)
        
        # 显示进度信息：种子、保存路径、耗时
        msg = f"✅ 第{i+1}/{batch_images}张完成，种子{seed+i}，耗时{img_time:.2f}秒 | 保存至: {filename}"
        print(msg)
        yield results, msg
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    
    # 生成完成后显示总结信息
    if results:
        total_time = time.time() - total_start_time
        avg_time = total_time / len(results) if results else 0
        msg = f"🎉 全部完成！共{len(results)}张，总耗时{total_time:.2f}秒，平均{avg_time:.2f}秒/张"
        print(msg)
        yield results, msg


def generate_t2i_image(
    prompt,
    negative_prompt,
    width, 
    height, 
    num_inference_steps, 
    batch_images, 
    seed_param, 
    lora_dropdown, 
    lora_weights,
    guidance_scale,
):
    global stop_generation, mode_loaded, lora_loaded, lora_loaded_weights
    if mode_loaded != "t2i_image" or lora_loaded != lora_dropdown or lora_loaded_weights != lora_weights:
        load_model("t2i_image", lora_dropdown, lora_weights)
        mode_loaded = "t2i_image"
        lora_loaded, lora_loaded_weights = lora_dropdown, lora_weights
        # 模型切换时清除缓存
        global prompt_embeds_cache
        prompt_embeds_cache = {"key": None, "prompt_embeds": None, "prompt_embeds_mask": None}
    results = []
    if seed_param < 0:
        seed = random.randint(0, np.iinfo(np.int32).max)
    else:
        seed = seed_param
    total_start_time = time.time()
    inference_times = []
    for i in range(batch_images):
        if stop_generation:
            stop_generation = False
            yield results, f"✅ 生成已中止，最后种子数{seed+i-1}"
            break
        img_start_time = time.time()
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"outputs/{timestamp}.png"
        generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed+i)
        with torch.inference_mode():
            output = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt if negative_prompt else None,
                height=height,
                width=width,
                cfg_normalization=False,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                generator=generator,
            )
        image = output.images[0]
        image.save(filename)
        results.append(image)
        
        # 计算单张图生成时间
        img_time = time.time() - img_start_time
        inference_times.append(img_time)
        
        # 显示进度信息：种子、保存路径、耗时
        msg = f"✅ 第{i+1}/{batch_images}张完成，种子{seed+i}，耗时{img_time:.2f}秒 | 保存至: {filename}"
        print(msg)
        yield results, msg
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    
    # 生成完成后显示总结信息
    if results:
        total_time = time.time() - total_start_time
        avg_time = total_time / len(results) if results else 0
        msg = f"🎉 全部完成！共{len(results)}张，总耗时{total_time:.2f}秒，平均{avg_time:.2f}秒/张"
        print(msg)
        yield results, msg


def generate_i2i(
    init_image,
    prompt, 
    width, 
    height, 
    strength,
    num_inference_steps, 
    batch_images, 
    seed_param, 
    lora_dropdown, 
    lora_weights,
):
    global stop_generation, mode_loaded, lora_loaded, lora_loaded_weights
    if mode_loaded != "i2i" or lora_loaded != lora_dropdown or lora_loaded_weights != lora_weights:
        load_model("i2i", lora_dropdown, lora_weights)
        mode_loaded = "i2i"
        lora_loaded, lora_loaded_weights = lora_dropdown, lora_weights
        # 模型切换时清除缓存
        global prompt_embeds_cache
        prompt_embeds_cache = {"key": None, "prompt_embeds": None, "prompt_embeds_mask": None}
    results = []
    if seed_param < 0:
        seed = random.randint(0, np.iinfo(np.int32).max)
    else:
        seed = seed_param
    prompt_embeds, _ = get_cached_prompt_embeds(prompt)
    if init_image.size != (width, height):
        init_image = init_image.resize((width, height))
    total_start_time = time.time()
    inference_times = []
    for i in range(batch_images):
        if stop_generation:
            stop_generation = False
            yield results, f"✅ 生成已中止，最后种子数{seed+i-1}"
            break
        img_start_time = time.time()
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"outputs/{timestamp}.png"
        with torch.inference_mode():
            output = pipe(
                image=init_image,
                prompt=prompt,
                height=height,
                width=width,
                strength=strength,
                num_inference_steps=num_inference_steps, 
                guidance_scale=0.0, 
                generator=torch.Generator().manual_seed(seed+i),
                prompt_embeds=prompt_embeds,
            )
        image = output.images[0]
        image.save(filename)
        results.append(image)
        
        # 计算单张图生成时间
        img_time = time.time() - img_start_time
        inference_times.append(img_time)
        
        # 显示进度信息：种子、保存路径、耗时
        msg = f"✅ 第{i+1}/{batch_images}张完成，种子{seed+i}，耗时{img_time:.2f}秒 | 保存至: {filename}"
        print(msg)
        yield results, msg
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    
    # 生成完成后显示总结信息
    if results:
        total_time = time.time() - total_start_time
        avg_time = total_time / len(results) if results else 0
        msg = f"🎉 全部完成！共{len(results)}张，总耗时{total_time:.2f}秒，平均{avg_time:.2f}秒/张"
        print(msg)
        yield results, msg


def generate_con(
    prompt, 
    width, 
    height, 
    num_inference_steps, 
    strength,
    batch_images, 
    seed_param, 
    lora_dropdown, 
    lora_weights,
    control_image=None,
    image=None,
):
    global stop_generation, mode_loaded, lora_loaded, lora_loaded_weights
    if mode_loaded != "con" or lora_loaded != lora_dropdown or lora_loaded_weights != lora_weights:
        load_model("con", lora_dropdown, lora_weights)
        mode_loaded = "con"
        lora_loaded, lora_loaded_weights = lora_dropdown, lora_weights
        # 模型切换时清除缓存
        global prompt_embeds_cache
        prompt_embeds_cache = {"key": None, "prompt_embeds": None, "prompt_embeds_mask": None}
    results = []
    if seed_param < 0:
        seed = random.randint(0, np.iinfo(np.int32).max)
    else:
        seed = seed_param
    prompt_embeds, _ = get_cached_prompt_embeds(prompt)
    if image["background"] is not None:
        # 处理蒙版图像
        mask_image = image["layers"][0]
        mask_image = mask_image .convert("RGBA")
        data = np.array(mask_image)
        # 修改蒙版颜色（黑色->白色，透明->黑色）
        black_pixels = (data[:, :, 0] == 0) & (data[:, :, 1] == 0) & (data[:, :, 2] == 0)
        data[black_pixels, :3] = [255, 255, 255]
        transparent_pixels = (data[:, :, 3] == 0)
        data[transparent_pixels, :3] = [0, 0, 0]
        mask_image = Image.fromarray(data)
        mask_image = get_image_latent(mask_image, sample_size=[height, width])[:, :1, 0]
        # 提取背景图像
        inpaint_image = load_image(image["background"])
        inpaint_image = get_image_latent(inpaint_image, sample_size=[height, width])[:, :, 0]
    else:
        inpaint_image = torch.zeros([1, 3, height, width])
        mask_image = torch.ones([1, 1, height, width]) * 255
    if control_image is not None:
        control_image = get_image_latent(control_image, sample_size=[height, width])[:, :, 0]
    else:
        control_image = torch.zeros([1, 3, height, width])
    total_start_time = time.time()
    inference_times = []
    for i in range(batch_images):
        if stop_generation:
            stop_generation = False
            yield results, f"✅ 生成已中止，最后种子数{seed+i-1}"
            break
        img_start_time = time.time()
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"outputs/{timestamp}.png"
        with torch.inference_mode():
            output = pipe(
                prompt=prompt,
                height=height,
                width=width,
                num_inference_steps=num_inference_steps, 
                guidance_scale=0.0, 
                generator=torch.Generator().manual_seed(seed+i),
                prompt_embeds=prompt_embeds,
                image=inpaint_image,
                mask_image=mask_image,
                control_image=control_image,
                control_context_scale=strength,
            )
        image = output.images[0]
        image.save(filename)
        results.append(image)
        
        # 计算单张图生成时间
        img_time = time.time() - img_start_time
        inference_times.append(img_time)
        
        # 显示进度信息：种子、保存路径、耗时
        msg = f"✅ 第{i+1}/{batch_images}张完成，种子{seed+i}，耗时{img_time:.2f}秒 | 保存至: {filename}"
        print(msg)
        yield results, msg
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    
    # 生成完成后显示总结信息
    if results:
        total_time = time.time() - total_start_time
        avg_time = total_time / len(results) if results else 0
        msg = f"🎉 全部完成！共{len(results)}张，总耗时{total_time:.2f}秒，平均{avg_time:.2f}秒/张"
        print(msg)
        yield results, msg
    

with gr.Blocks(title="Z-Image-diffusers", theme=gr.themes.Soft(font=[gr.themes.GoogleFont("IBM Plex Sans")])) as demo:
    gr.Markdown("""
            <div>
                <h2 style="font-size: 30px;text-align: center;">Z-Image-diffusers</h2>
            </div>
            <div style="text-align: center;">
                十字鱼
                <a href="https://space.bilibili.com/893892">🌐bilibili</a> 
                |Z-Image-diffusers
                <a href="https://github.com/gluttony-10/Z-Image-diffusers">🌐github</a> 
            </div>
            <div style="text-align: center; font-weight: bold; color: red;">
                ⚠️ 该演示仅供学术研究和体验使用。
            </div>
            """)
    with gr.Accordion("LoRA设置", open=False):
        with gr.Column():
            with gr.Row():
                lora_dropdown = gr.Dropdown(label="LoRA模型", info="存放LoRA模型到models/lora，可多选", choices=lora_choices, multiselect=True)
                lora_weights = gr.Textbox(label="LoRA权重", info="Lora权重，多个权重请用英文逗号隔开。例如：0.8,0.5,0.2", value="")
    with gr.Tabs():
        with gr.TabItem("文生图（Turbo）"):
            with gr.Row():
                with gr.Column():
                    prompt_t2i = gr.Textbox(label="提示词", placeholder="请输入提示词...")
                    generate_button_t2i = gr.Button("🖼️ 开始生成", variant='primary', scale=4)
                    with gr.Accordion("参数设置", open=True):
                        with gr.Row():
                            width_t2i = gr.Slider(label="宽度", minimum=256, maximum=3072, step=16, value=1024)
                            height_t2i = gr.Slider(label="高度", minimum=256, maximum=3072, step=16, value=1024)
                        with gr.Row():
                            exchange_button_t2i = gr.Button("🔄 交换宽高")
                            scale_1_5_button_t2i = gr.Button("1.5倍分辨率")
                        batch_images_t2i = gr.Slider(label="批量生成", minimum=1, maximum=100, step=1, value=1)
                        num_inference_steps_t2i = gr.Slider(label="采样步数（推荐9步）", minimum=1, maximum=100, step=1, value=9)
                        seed_param_t2i = gr.Number(label="种子，请输入自然数，-1为随机", value=-1)
                with gr.Column():
                    info_t2i = gr.Textbox(label="提示信息", interactive=False)
                    image_output_t2i = gr.Gallery(label="生成结果", interactive=False)
                    stop_button_t2i = gr.Button("中止生成", variant="stop")
        with gr.TabItem("文生图（Image）"):
            with gr.Row():
                with gr.Column():
                    prompt_t2i_image = gr.Textbox(label="提示词", placeholder="请输入提示词...")
                    negative_prompt_t2i_image = gr.Textbox(label="负面提示词", placeholder="请输入负面提示词...", value="")
                    generate_button_t2i_image = gr.Button("🖼️ 开始生成", variant='primary', scale=4)
                    with gr.Accordion("参数设置", open=True):
                        with gr.Row():
                            width_t2i_image = gr.Slider(label="宽度", minimum=256, maximum=3072, step=16, value=1024)
                            height_t2i_image = gr.Slider(label="高度", minimum=256, maximum=3072, step=16, value=1024)
                        with gr.Row():
                            exchange_button_t2i_image = gr.Button("🔄 交换宽高")
                            scale_1_5_button_t2i_image = gr.Button("1.5倍分辨率")
                        batch_images_t2i_image = gr.Slider(label="批量生成", minimum=1, maximum=100, step=1, value=1)
                        num_inference_steps_t2i_image = gr.Slider(label="采样步数（推荐50步）", minimum=1, maximum=100, step=1, value=50)
                        guidance_scale_t2i_image = gr.Slider(label="guidance_scale", minimum=0, maximum=10, step=0.1, value=4)
                        seed_param_t2i_image = gr.Number(label="种子，请输入自然数，-1为随机", value=-1)
                with gr.Column():
                    info_t2i_image = gr.Textbox(label="提示信息", interactive=False)
                    image_output_t2i_image = gr.Gallery(label="生成结果", interactive=False)
                    stop_button_t2i_image = gr.Button("中止生成", variant="stop")
        """with gr.TabItem("图生图"):
            with gr.Row():
                with gr.Column():
                    image_i2i = gr.Image(label="输入图片", type="pil", height=300)
                    prompt_i2i = gr.Textbox(label="提示词", placeholder="请输入提示词...")
                    generate_button_i2i = gr.Button("🖼️ 开始生成", variant='primary', scale=4)
                    with gr.Accordion("参数设置", open=True):
                        with gr.Row():
                            width_i2i = gr.Slider(label="宽度", minimum=256, maximum=3072, step=16, value=1024)
                            height_i2i = gr.Slider(label="高度", minimum=256, maximum=3072, step=16, value=1024)
                        with gr.Row():
                            exchange_button_i2i = gr.Button("🔄 交换宽高")
                            scale_1_5_button_i2i = gr.Button("1.5倍分辨率")
                        strength_i2i = gr.Slider(label="strength（推荐0.6）", minimum=0, maximum=1, step=0.01, value=0.6)
                        batch_images_i2i = gr.Slider(label="批量生成", minimum=1, maximum=100, step=1, value=1)
                        num_inference_steps_i2i = gr.Slider(label="采样步数（推荐9步）", minimum=1, maximum=100, step=1, value=9)
                        seed_param_i2i = gr.Number(label="种子，请输入自然数，-1为随机", value=-1)
                with gr.Column():
                    info_i2i = gr.Textbox(label="提示信息", interactive=False)
                    image_output_i2i = gr.Gallery(label="生成结果", interactive=False)
                    stop_button_i2i = gr.Button("中止生成", variant="stop")"""
        """with gr.TabItem("ControlNet"):
            with gr.Row():
                with gr.Column():
                    image_con = gr.Image(label="输入控制图片（支持Canny、HED、Depth、Pose、MLSD）", type="pil", height=300)
                    prompt_con = gr.Textbox(label="提示词", placeholder="请输入提示词...")
                    generate_button_con = gr.Button("🖼️ 开始生成", variant='primary', scale=4)
                    with gr.Accordion("参数设置", open=True):
                        with gr.Row():
                            width_con = gr.Slider(label="宽度", minimum=256, maximum=3072, step=16, value=1024)
                            height_con = gr.Slider(label="高度", minimum=256, maximum=3072, step=16, value=1024)
                        with gr.Row():
                            exchange_button_con = gr.Button("🔄 交换宽高")
                            scale_1_5_button_con = gr.Button("1.5倍分辨率")
                        strength_con = gr.Slider(label="strength（推荐0.75）", minimum=0, maximum=1, step=0.01, value=0.75)
                        batch_images_con = gr.Slider(label="批量生成", minimum=1, maximum=100, step=1, value=1)
                        num_inference_steps_con = gr.Slider(label="采样步数（推荐9步）", minimum=1, maximum=100, step=1, value=9)
                        seed_param_con = gr.Number(label="种子，请输入自然数，-1为随机", value=-1)
                with gr.Column():
                    info_con = gr.Textbox(label="提示信息", interactive=False)
                    image_output_con = gr.Gallery(label="生成结果", interactive=False)
                    stop_button_con = gr.Button("中止生成", variant="stop")"""
        with gr.TabItem("ControlNet&Inpaint"):
            with gr.Row():
                with gr.Column():
                    image_inp = gr.ImageMask(label="输入修改图片和蒙版", type="pil", height=400)
                    image_coni = gr.Image(label="输入控制图片（支持Canny、HED、Depth、Pose、MLSD）", type="pil", height=300)
                    prompt_coni = gr.Textbox(label="提示词", placeholder="请输入提示词...")
                    generate_button_coni = gr.Button("🖼️ 开始生成", variant='primary', scale=4)
                    with gr.Accordion("参数设置", open=True):
                        with gr.Row():
                            width_coni = gr.Slider(label="宽度", minimum=256, maximum=3072, step=16, value=1024)
                            height_coni = gr.Slider(label="高度", minimum=256, maximum=3072, step=16, value=1024)
                        with gr.Row():
                            exchange_button_coni = gr.Button("🔄 交换宽高")
                            scale_1_5_button_coni = gr.Button("1.5倍分辨率")
                        strength_coni = gr.Slider(label="strength（推荐0.75）", minimum=0, maximum=1, step=0.01, value=0.75)
                        batch_images_coni = gr.Slider(label="批量生成", minimum=1, maximum=100, step=1, value=1)
                        num_inference_steps_coni = gr.Slider(label="采样步数（推荐9步）", minimum=1, maximum=100, step=1, value=9)
                        seed_param_coni = gr.Number(label="种子，请输入自然数，-1为随机", value=-1)
                with gr.Column():
                    info_coni = gr.Textbox(label="提示信息", interactive=False)
                    image_output_coni = gr.Gallery(label="生成结果", interactive=False)
                    stop_button_coni = gr.Button("中止生成", variant="stop")
    # 文生图（Turbo）
    gr.on(
        triggers=[generate_button_t2i.click, prompt_t2i.submit],
        fn = generate_t2i,
        inputs = [
            prompt_t2i,
            width_t2i,
            height_t2i,
            num_inference_steps_t2i,
            batch_images_t2i,
            seed_param_t2i,
            lora_dropdown, 
            lora_weights,
        ],
        outputs = [image_output_t2i, info_t2i]
    )
    exchange_button_t2i.click(
        fn=exchange_width_height, 
        inputs=[width_t2i, height_t2i], 
        outputs=[width_t2i, height_t2i, info_t2i]
    )
    scale_1_5_button_t2i.click(
        fn=scale_resolution_1_5,
        inputs=[width_t2i, height_t2i],
        outputs=[width_t2i, height_t2i, info_t2i]
    )
    stop_button_t2i.click(
        fn=stop_generate, 
        inputs=[], 
        outputs=[info_t2i]
    )
    # 文生图（Image）
    gr.on(
        triggers=[generate_button_t2i_image.click, prompt_t2i_image.submit, negative_prompt_t2i_image.submit],
        fn = generate_t2i_image,
        inputs = [
            prompt_t2i_image,
            negative_prompt_t2i_image,
            width_t2i_image,
            height_t2i_image,
            num_inference_steps_t2i_image,
            batch_images_t2i_image,
            seed_param_t2i_image,
            lora_dropdown, 
            lora_weights,
            guidance_scale_t2i_image,
        ],
        outputs = [image_output_t2i_image, info_t2i_image]
    )
    exchange_button_t2i_image.click(
        fn=exchange_width_height, 
        inputs=[width_t2i_image, height_t2i_image], 
        outputs=[width_t2i_image, height_t2i_image, info_t2i_image]
    )
    scale_1_5_button_t2i_image.click(
        fn=scale_resolution_1_5,
        inputs=[width_t2i_image, height_t2i_image],
        outputs=[width_t2i_image, height_t2i_image, info_t2i_image]
    )
    stop_button_t2i_image.click(
        fn=stop_generate, 
        inputs=[], 
        outputs=[info_t2i_image]
    )
    # 图生图  
    """gr.on(
        triggers=[generate_button_i2i.click, prompt_i2i.submit],
        fn = generate_i2i,
        inputs = [
            image_i2i,
            prompt_i2i,
            width_i2i,
            height_i2i,
            strength_i2i,
            num_inference_steps_i2i,
            batch_images_i2i,
            seed_param_i2i,
            lora_dropdown, 
            lora_weights,
        ],
        outputs = [image_output_i2i, info_i2i]
    )
    exchange_button_i2i.click(
        fn=exchange_width_height, 
        inputs=[width_i2i, height_i2i], 
        outputs=[width_i2i, height_i2i, info_i2i]
    )
    scale_1_5_button_i2i.click(
        fn=scale_resolution_1_5,
        inputs=[width_i2i, height_i2i],
        outputs=[width_i2i, height_i2i, info_i2i]
    )
    stop_button_i2i.click(
        fn=stop_generate, 
        inputs=[], 
        outputs=[info_i2i]
    )
    image_i2i.upload(
        fn=adjust_width_height, 
        inputs=[image_i2i], 
        outputs=[width_i2i, height_i2i, info_i2i]
    )"""
    # ControlNet
    """gr.on(
        triggers=[generate_button_con.click, prompt_con.submit],
        fn = generate_con,
        inputs = [
            prompt_con,
            width_con,
            height_con,
            num_inference_steps_con,
            strength_con,
            batch_images_con,
            seed_param_con,
            lora_dropdown, 
            lora_weights,
            image_con,
        ],
        outputs = [image_output_con, info_con]
    )
    exchange_button_con.click(
        fn=exchange_width_height, 
        inputs=[width_con, height_con], 
        outputs=[width_con, height_con, info_con]
    )
    scale_1_5_button_con.click(
        fn=scale_resolution_1_5,
        inputs=[width_con, height_con],
        outputs=[width_con, height_con, info_con]
    )
    stop_button_con.click(
        fn=stop_generate, 
        inputs=[], 
        outputs=[info_con]
    )
    image_con.upload(
        fn=adjust_width_height, 
        inputs=[image_con], 
        outputs=[width_con, height_con, info_con]
    )"""
    # ControlNet-Inpaint
    gr.on(
        triggers=[generate_button_coni.click, prompt_coni.submit],
        fn = generate_con,
        inputs = [
            prompt_coni,
            width_coni,
            height_coni,
            num_inference_steps_coni,
            strength_coni,
            batch_images_coni,
            seed_param_coni,
            lora_dropdown, 
            lora_weights,
            image_coni,
            image_inp,
        ],
        outputs = [image_output_coni, info_coni]
    )
    exchange_button_coni.click(
        fn=exchange_width_height, 
        inputs=[width_coni, height_coni], 
        outputs=[width_coni, height_coni, info_coni]
    )
    scale_1_5_button_coni.click(
        fn=scale_resolution_1_5,
        inputs=[width_coni, height_coni],
        outputs=[width_coni, height_coni, info_coni]
    )
    stop_button_coni.click(
        fn=stop_generate, 
        inputs=[], 
        outputs=[info_coni]
    )
    image_coni.upload(
        fn=adjust_width_height, 
        inputs=[image_coni], 
        outputs=[width_coni, height_coni, info_coni]
    )

if __name__ == "__main__": 
    demo.launch(
        server_name=args.server_name, 
        server_port=find_port(args.server_port),
        share=args.share, 
        mcp_server=args.mcp_server,
        inbrowser=True,
    )