Add files using upload-large-folder tool
Browse files- .gitattributes +12 -0
- README.md +3 -1
- Z-Image-Launcher.exe +2 -2
- app.py +1040 -0
- assets/DMDR.webp +3 -0
- assets/Z-Image-Gallery.pdf +3 -0
- assets/architecture.webp +3 -0
- assets/decoupled-dmd.webp +3 -0
- assets/image_arena_all.jpg +3 -0
- assets/image_arena_os.jpg +3 -0
- assets/leaderboard.png +3 -0
- assets/reasoning.png +3 -0
- assets/showcase.jpg +3 -0
- assets/showcase_editing.png +3 -0
- assets/showcase_realistic.png +3 -0
- assets/showcase_rendering.png +3 -0
- assets/training_pipeline.jpg +3 -0
.gitattributes
CHANGED
|
@@ -585,3 +585,15 @@ python_env/lib/site-packages/pandas/tests/indexing/__pycache__/test_loc.cpython-
|
|
| 585 |
python_env/lib/site-packages/pandas/tests/io/__pycache__/test_sql.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text
|
| 586 |
python_env/lib/site-packages/pandas/tests/tools/__pycache__/test_to_datetime.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text
|
| 587 |
WebView2Loader.dll filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 585 |
python_env/lib/site-packages/pandas/tests/io/__pycache__/test_sql.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text
|
| 586 |
python_env/lib/site-packages/pandas/tests/tools/__pycache__/test_to_datetime.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text
|
| 587 |
WebView2Loader.dll filter=lfs diff=lfs merge=lfs -text
|
| 588 |
+
assets/architecture.webp filter=lfs diff=lfs merge=lfs -text
|
| 589 |
+
assets/DMDR.webp filter=lfs diff=lfs merge=lfs -text
|
| 590 |
+
assets/image_arena_all.jpg filter=lfs diff=lfs merge=lfs -text
|
| 591 |
+
assets/decoupled-dmd.webp filter=lfs diff=lfs merge=lfs -text
|
| 592 |
+
assets/leaderboard.png filter=lfs diff=lfs merge=lfs -text
|
| 593 |
+
assets/image_arena_os.jpg filter=lfs diff=lfs merge=lfs -text
|
| 594 |
+
assets/showcase.jpg filter=lfs diff=lfs merge=lfs -text
|
| 595 |
+
assets/training_pipeline.jpg filter=lfs diff=lfs merge=lfs -text
|
| 596 |
+
assets/reasoning.png filter=lfs diff=lfs merge=lfs -text
|
| 597 |
+
assets/showcase_editing.png filter=lfs diff=lfs merge=lfs -text
|
| 598 |
+
assets/showcase_rendering.png filter=lfs diff=lfs merge=lfs -text
|
| 599 |
+
assets/Z-Image-Gallery.pdf filter=lfs diff=lfs merge=lfs -text
|
README.md
CHANGED
|
@@ -12,7 +12,9 @@
|
|
| 12 |
<a href="https://arxiv.org/abs/2511.22699" target="_blank"><img src="https://img.shields.io/badge/Report-b5212f.svg?logo=arxiv" height="21px"></a>
|
| 13 |
|
| 14 |
|
| 15 |
-
|
|
|
|
|
|
|
| 16 |
|
| 17 |
</div>
|
| 18 |
|
|
|
|
| 12 |
<a href="https://arxiv.org/abs/2511.22699" target="_blank"><img src="https://img.shields.io/badge/Report-b5212f.svg?logo=arxiv" height="21px"></a>
|
| 13 |
|
| 14 |
|
| 15 |
+
这是一个基于通义Z-Image-Turbo 的项目,主要是为了满足低显存的设备能够快速出图,同时也增加了一些常用功能,图生图,局部重绘之类的。
|
| 16 |
+
|
| 17 |
+
This is a project based on Tongyi Z-Image-Turbo. It is mainly designed to enable devices with low video memory to generate images quickly. Meanwhile, some common functions such as image-to-image generation and local redrawing have also been added.
|
| 18 |
|
| 19 |
</div>
|
| 20 |
|
Z-Image-Launcher.exe
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:d5673674af1ea04b8424396da60676f91d9e3b4690d61f071fafd9b0705a76be
|
| 3 |
+
size 1413433
|
app.py
ADDED
|
@@ -0,0 +1,1040 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import time # 【新增】引入time模块用于退出延迟
|
| 3 |
+
|
| 4 |
+
# 必须处于文件最顶端:环境配置
|
| 5 |
+
os.environ["DIFFUSERS_USE_PEFT_BACKEND"] = "1"
|
| 6 |
+
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
|
| 7 |
+
os.environ["HF_HUB_OFFLINE"] = "1"
|
| 8 |
+
|
| 9 |
+
import sys
|
| 10 |
+
import torch
|
| 11 |
+
import psutil
|
| 12 |
+
import random
|
| 13 |
+
import re
|
| 14 |
+
import uuid
|
| 15 |
+
import gc
|
| 16 |
+
from datetime import datetime
|
| 17 |
+
from PIL import Image, ImageFilter, ImageOps, ImageEnhance
|
| 18 |
+
|
| 19 |
+
# 配置基础路径
|
| 20 |
+
current_dir = os.path.dirname(os.path.abspath(__file__))
|
| 21 |
+
if current_dir not in sys.path:
|
| 22 |
+
sys.path.append(current_dir)
|
| 23 |
+
|
| 24 |
+
# 目录配置
|
| 25 |
+
DEFAULT_MODEL_PATH = os.path.join(current_dir, "ckpts", "Z-Image-Turbo")
|
| 26 |
+
LORA_ROOT = os.path.join(current_dir, "lora")
|
| 27 |
+
OUTPUT_ROOT = os.path.join(current_dir, "outputs")
|
| 28 |
+
MOD_VAE_DIR = os.path.join(current_dir, "Mod", "vae")
|
| 29 |
+
MOD_TRANS_DIR = os.path.join(current_dir, "Mod", "transformer")
|
| 30 |
+
for p in [LORA_ROOT, OUTPUT_ROOT, MOD_VAE_DIR, MOD_TRANS_DIR]:
|
| 31 |
+
os.makedirs(p, exist_ok=True)
|
| 32 |
+
|
| 33 |
+
try:
|
| 34 |
+
import gradio as gr
|
| 35 |
+
from diffusers import ZImagePipeline, ZImageImg2ImgPipeline, AutoencoderKL
|
| 36 |
+
from safetensors.torch import load_file
|
| 37 |
+
except ImportError as e:
|
| 38 |
+
print(f"❌ 核心库导入失败: {e}")
|
| 39 |
+
sys.exit(1)
|
| 40 |
+
|
| 41 |
+
# ==========================================
|
| 42 |
+
# 设备探测与硬件报告
|
| 43 |
+
# ==========================================
|
| 44 |
+
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
| 45 |
+
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
|
| 46 |
+
is_interrupted = False
|
| 47 |
+
|
| 48 |
+
print("\n" + "="*50)
|
| 49 |
+
if DEVICE == "cuda":
|
| 50 |
+
GPU_NAME = torch.cuda.get_device_name(0)
|
| 51 |
+
TOTAL_VRAM = torch.cuda.get_device_properties(0).total_memory
|
| 52 |
+
print(f"✅ 运行模式: [ GPU ]")
|
| 53 |
+
print(f"核心型号: {GPU_NAME}")
|
| 54 |
+
print(f"显存总量: {TOTAL_VRAM/1024**3:.2f} GB")
|
| 55 |
+
else:
|
| 56 |
+
TOTAL_VRAM = 0
|
| 57 |
+
print(f"⚠️ 运行模式: [ CPU ]")
|
| 58 |
+
print("="*50 + "\n")
|
| 59 |
+
|
| 60 |
+
# ==========================================
|
| 61 |
+
# 显存与工具函数
|
| 62 |
+
# ==========================================
|
| 63 |
+
def get_vram_info():
|
| 64 |
+
if DEVICE == "cuda":
|
| 65 |
+
reserved = torch.cuda.memory_reserved(0)
|
| 66 |
+
allocated = torch.cuda.memory_allocated(0)
|
| 67 |
+
usage_pct = (reserved / TOTAL_VRAM) * 100 if TOTAL_VRAM > 0 else 0
|
| 68 |
+
vram_str = (
|
| 69 |
+
f"显存占用: {usage_pct:.1f}% "
|
| 70 |
+
f"({reserved/1024**3:.2f}GB / {TOTAL_VRAM/1024**3:.2f}GB)"
|
| 71 |
+
)
|
| 72 |
+
else:
|
| 73 |
+
usage_pct = 0
|
| 74 |
+
vram_str = "显存占用: CPU 模式"
|
| 75 |
+
|
| 76 |
+
mem = psutil.virtual_memory()
|
| 77 |
+
ram_str = (
|
| 78 |
+
f"内存占用: {mem.percent:.1f}% "
|
| 79 |
+
f"({(mem.total - mem.available)/1024**3:.2f}GB / {mem.total/1024**3:.2f}GB)"
|
| 80 |
+
)
|
| 81 |
+
status = f"{vram_str} | {ram_str}"
|
| 82 |
+
return usage_pct, status
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def auto_flush_vram(threshold=90):
|
| 86 |
+
usage_pct, _ = get_vram_info()
|
| 87 |
+
if usage_pct > threshold:
|
| 88 |
+
gc.collect()
|
| 89 |
+
torch.cuda.empty_cache()
|
| 90 |
+
return True
|
| 91 |
+
return False
|
| 92 |
+
|
| 93 |
+
def scan_lora_files():
|
| 94 |
+
if not os.path.exists(LORA_ROOT): return []
|
| 95 |
+
return sorted([f for f in os.listdir(LORA_ROOT) if f.lower().endswith(".safetensors")])
|
| 96 |
+
|
| 97 |
+
def scan_model_items(base_path):
|
| 98 |
+
if not os.path.exists(base_path): return []
|
| 99 |
+
items = []
|
| 100 |
+
for f in os.listdir(base_path):
|
| 101 |
+
full_path = os.path.join(base_path, f)
|
| 102 |
+
if os.path.isdir(full_path):
|
| 103 |
+
items.append(f)
|
| 104 |
+
elif f.lower().endswith((".safetensors", ".bin", ".pt")):
|
| 105 |
+
items.append(f)
|
| 106 |
+
return sorted(items)
|
| 107 |
+
|
| 108 |
+
# ==========================================
|
| 109 |
+
# 全局 LoRA 文件列表 (启动时扫描)
|
| 110 |
+
# ==========================================
|
| 111 |
+
LORA_FILES = scan_lora_files()
|
| 112 |
+
print(f"🔍 已检测到 {len(LORA_FILES)} 个 LoRA 文件,正在生成独立控件...")
|
| 113 |
+
if len(LORA_FILES) > 30:
|
| 114 |
+
print("⚠️ 警告: LoRA 数量较多,生成界面可能需要几秒钟...")
|
| 115 |
+
|
| 116 |
+
# ==========================================
|
| 117 |
+
# 模型管理器 (修改版:支持性能模式切换)
|
| 118 |
+
# ==========================================
|
| 119 |
+
class ModelManager:
|
| 120 |
+
def __init__(self):
|
| 121 |
+
self.pipe = None
|
| 122 |
+
self.current_state = {
|
| 123 |
+
"mode": None,
|
| 124 |
+
"t_choice": None,
|
| 125 |
+
"v_choice": None,
|
| 126 |
+
"perf_mode": None # 【新增】记录当前性能模式
|
| 127 |
+
}
|
| 128 |
+
self.current_loras = []
|
| 129 |
+
self.current_weights_map = {}
|
| 130 |
+
|
| 131 |
+
def _clear_pipeline(self):
|
| 132 |
+
if self.pipe is not None:
|
| 133 |
+
print(f"🧹 正在销毁旧管道以释放显存...")
|
| 134 |
+
try:
|
| 135 |
+
self.pipe.unload_lora_weights()
|
| 136 |
+
except:
|
| 137 |
+
pass
|
| 138 |
+
del self.pipe
|
| 139 |
+
self.pipe = None
|
| 140 |
+
if hasattr(sys, 'last_traceback'):
|
| 141 |
+
del sys.last_traceback
|
| 142 |
+
for _ in range(3):
|
| 143 |
+
gc.collect()
|
| 144 |
+
torch.cuda.empty_cache()
|
| 145 |
+
torch.cuda.ipc_collect()
|
| 146 |
+
if DEVICE == "cuda":
|
| 147 |
+
res = torch.cuda.memory_reserved(0) / 1024**3
|
| 148 |
+
print(f"✨ 显存已深度清理,当前占用: {res:.2f} GB")
|
| 149 |
+
|
| 150 |
+
def _init_pipeline_base(self, mode):
|
| 151 |
+
if mode == 'txt':
|
| 152 |
+
print("🚀 初始化基础 Pipeline (文生图)...")
|
| 153 |
+
return ZImagePipeline.from_pretrained(DEFAULT_MODEL_PATH, torch_dtype=DTYPE, local_files_only=True)
|
| 154 |
+
else:
|
| 155 |
+
print("🚀 初始化基础 Pipeline (图生图)...")
|
| 156 |
+
return ZImageImg2ImgPipeline.from_pretrained(DEFAULT_MODEL_PATH, torch_dtype=DTYPE, local_files_only=True)
|
| 157 |
+
|
| 158 |
+
def _inject_components(self, pipe, t_choice, v_choice):
|
| 159 |
+
if t_choice != "default":
|
| 160 |
+
t_path = os.path.join(MOD_TRANS_DIR, t_choice)
|
| 161 |
+
if os.path.isfile(t_path):
|
| 162 |
+
print(f"📦 载入 Transformer: {t_choice}")
|
| 163 |
+
state_dict = load_file(t_path, device="cpu")
|
| 164 |
+
processed = {}
|
| 165 |
+
prefix = "model.diffusion_model."
|
| 166 |
+
for k, v in state_dict.items():
|
| 167 |
+
new_k = k[len(prefix):] if k.startswith(prefix) else k
|
| 168 |
+
processed[new_k] = v.to(DTYPE)
|
| 169 |
+
pipe.transformer.load_state_dict(processed, strict=False, assign=True)
|
| 170 |
+
del state_dict, processed, v
|
| 171 |
+
gc.collect()
|
| 172 |
+
|
| 173 |
+
if v_choice != "default":
|
| 174 |
+
vae_path = os.path.join(MOD_VAE_DIR, v_choice)
|
| 175 |
+
print(f"📦 载入 VAE: {v_choice}")
|
| 176 |
+
if os.path.isfile(vae_path):
|
| 177 |
+
pipe.vae = AutoencoderKL.from_single_file(vae_path, torch_dtype=DTYPE)
|
| 178 |
+
else:
|
| 179 |
+
pipe.vae = AutoencoderKL.from_pretrained(vae_path, torch_dtype=DTYPE)
|
| 180 |
+
return pipe
|
| 181 |
+
|
| 182 |
+
def _apply_loras(self, pipe, selected_loras, weights_map):
|
| 183 |
+
if self.current_loras == selected_loras and self.current_weights_map == weights_map:
|
| 184 |
+
return
|
| 185 |
+
|
| 186 |
+
print("🎸 正在配置 LoRA (独立权重模式)...")
|
| 187 |
+
try:
|
| 188 |
+
pipe.unload_lora_weights()
|
| 189 |
+
except Exception:
|
| 190 |
+
pass
|
| 191 |
+
|
| 192 |
+
if not selected_loras:
|
| 193 |
+
self.current_loras = []
|
| 194 |
+
self.current_weights_map = {}
|
| 195 |
+
return
|
| 196 |
+
|
| 197 |
+
active_adapters = []
|
| 198 |
+
adapter_weights = []
|
| 199 |
+
|
| 200 |
+
for lora_file in selected_loras:
|
| 201 |
+
adapter_name = re.sub(r"[^a-zA-Z0-9_]", "_", os.path.splitext(lora_file)[0])
|
| 202 |
+
weight = weights_map.get(lora_file, 1.0)
|
| 203 |
+
|
| 204 |
+
try:
|
| 205 |
+
pipe.load_lora_weights(LORA_ROOT, weight_name=lora_file, adapter_name=adapter_name)
|
| 206 |
+
active_adapters.append(adapter_name)
|
| 207 |
+
adapter_weights.append(weight)
|
| 208 |
+
except Exception as e:
|
| 209 |
+
print(f"⚠️ LoRA {lora_file} 加载失败: {e}")
|
| 210 |
+
|
| 211 |
+
if active_adapters:
|
| 212 |
+
pipe.set_adapters(active_adapters, adapter_weights=adapter_weights)
|
| 213 |
+
|
| 214 |
+
self.current_loras = list(selected_loras)
|
| 215 |
+
self.current_weights_map = dict(weights_map)
|
| 216 |
+
|
| 217 |
+
def get_pipeline(self, t_choice, v_choice, selected_loras, weights_map, mode='txt', perf_mode="高端机 (显存>=20GB)"):
|
| 218 |
+
# 判断是否为低端机模式
|
| 219 |
+
is_low_vram = (perf_mode == "低端机 (显存优化)")
|
| 220 |
+
|
| 221 |
+
# 检查是否需要重建模型 (包含性能模式变更的检查)
|
| 222 |
+
need_rebuild = (
|
| 223 |
+
self.pipe is None or
|
| 224 |
+
self.current_state["mode"] != mode or
|
| 225 |
+
self.current_state["t_choice"] != t_choice or
|
| 226 |
+
self.current_state["v_choice"] != v_choice or
|
| 227 |
+
self.current_state["perf_mode"] != perf_mode # 【新增】性能模式变更需重建
|
| 228 |
+
)
|
| 229 |
+
|
| 230 |
+
if need_rebuild:
|
| 231 |
+
self._clear_pipeline()
|
| 232 |
+
try:
|
| 233 |
+
temp_pipe = self._init_pipeline_base(mode)
|
| 234 |
+
temp_pipe = self._inject_components(temp_pipe, t_choice, v_choice)
|
| 235 |
+
|
| 236 |
+
# 【核心修改】根据选择应用不同的显存优化策略
|
| 237 |
+
if DEVICE == "cuda":
|
| 238 |
+
if is_low_vram:
|
| 239 |
+
temp_pipe.enable_sequential_cpu_offload()
|
| 240 |
+
print(" [System] 已启用低显存优化模式")
|
| 241 |
+
else:
|
| 242 |
+
temp_pipe.to("cuda")
|
| 243 |
+
print(" [System] 已启用高端机模式")
|
| 244 |
+
|
| 245 |
+
self.pipe = temp_pipe
|
| 246 |
+
|
| 247 |
+
self.current_state = {
|
| 248 |
+
"mode": mode,
|
| 249 |
+
"t_choice": t_choice,
|
| 250 |
+
"v_choice": v_choice,
|
| 251 |
+
"perf_mode": perf_mode
|
| 252 |
+
}
|
| 253 |
+
self.current_loras = []
|
| 254 |
+
self.current_weights_map = {}
|
| 255 |
+
|
| 256 |
+
except Exception as e:
|
| 257 |
+
self._clear_pipeline()
|
| 258 |
+
raise gr.Error(f"模型加载崩溃: {str(e)}\n请检查显存或模型文件。")
|
| 259 |
+
|
| 260 |
+
self._apply_loras(self.pipe, selected_loras, weights_map)
|
| 261 |
+
return self.pipe
|
| 262 |
+
|
| 263 |
+
manager = ModelManager()
|
| 264 |
+
|
| 265 |
+
# ==========================================
|
| 266 |
+
# 进度回调
|
| 267 |
+
# ==========================================
|
| 268 |
+
def make_progress_callback(progress, total_steps, refresh_interval=2):
|
| 269 |
+
def _callback(pipe, step, timestep, callback_kwargs):
|
| 270 |
+
global is_interrupted
|
| 271 |
+
if is_interrupted:
|
| 272 |
+
raise gr.Error("🛑 任务已手动停止")
|
| 273 |
+
step_idx = step + 1
|
| 274 |
+
frac = step_idx / total_steps
|
| 275 |
+
status_suffix = ""
|
| 276 |
+
if step_idx % refresh_interval == 0 or step_idx == total_steps:
|
| 277 |
+
_, mem_status = get_vram_info()
|
| 278 |
+
status_suffix = f"\n{mem_status}"
|
| 279 |
+
progress(frac, desc=f"Diffusion Step {step_idx}/{total_steps}{status_suffix}")
|
| 280 |
+
return callback_kwargs
|
| 281 |
+
return _callback
|
| 282 |
+
|
| 283 |
+
# ==========================================
|
| 284 |
+
# 核心逻辑 (解析独立控件传入的参数)
|
| 285 |
+
# ==========================================
|
| 286 |
+
def process_lora_inputs(lora_checks, lora_weights):
|
| 287 |
+
selected = []
|
| 288 |
+
weights_map = {}
|
| 289 |
+
for i, fname in enumerate(LORA_FILES):
|
| 290 |
+
if i < len(lora_checks) and lora_checks[i]:
|
| 291 |
+
selected.append(fname)
|
| 292 |
+
if i < len(lora_weights):
|
| 293 |
+
weights_map[fname] = lora_weights[i]
|
| 294 |
+
else:
|
| 295 |
+
weights_map[fname] = 1.0
|
| 296 |
+
return selected, weights_map
|
| 297 |
+
|
| 298 |
+
# 【新增】LoRA 刷新函数
|
| 299 |
+
def refresh_lora_list():
|
| 300 |
+
global LORA_FILES
|
| 301 |
+
LORA_FILES = scan_lora_files()
|
| 302 |
+
count = len(LORA_FILES)
|
| 303 |
+
print(f"🔄 LoRA 刷新完成,当前检测到 {count} 个文件。")
|
| 304 |
+
# 提示用户刷新网页以看到新控件(因为 Gradio 静态布局限制)
|
| 305 |
+
msg = f"✅ 扫描完成!检测到 **{count}** 个 LoRA 文件。\n\n*(注:若新增了文件,请**刷新浏览器页面 (F5)** 以加载新的勾选框和滑块)*"
|
| 306 |
+
return gr.update(value=msg)
|
| 307 |
+
|
| 308 |
+
# 【新增】更新 Prompt UI 的辅助函数
|
| 309 |
+
def update_prompt_ui_base(prompt, *lora_ui_args):
|
| 310 |
+
num_loras = len(LORA_FILES)
|
| 311 |
+
if num_loras == 0:
|
| 312 |
+
return prompt
|
| 313 |
+
|
| 314 |
+
checks = lora_ui_args[:num_loras]
|
| 315 |
+
weights = lora_ui_args[num_loras:num_loras*2]
|
| 316 |
+
|
| 317 |
+
clean_p = re.sub(r"\s*<lora:[^>]+>", "", prompt or "").strip()
|
| 318 |
+
|
| 319 |
+
new_tags = []
|
| 320 |
+
for i, fname in enumerate(LORA_FILES):
|
| 321 |
+
if i < len(checks) and checks[i]:
|
| 322 |
+
w = weights[i] if i < len(weights) else 1.0
|
| 323 |
+
name = os.path.splitext(fname)[0]
|
| 324 |
+
alpha_str = f"{w:.2f}".rstrip("0").rstrip(".")
|
| 325 |
+
new_tags.append(f"<lora:{name}:{alpha_str}>")
|
| 326 |
+
|
| 327 |
+
if new_tags:
|
| 328 |
+
return f"{clean_p} {' '.join(new_tags)}"
|
| 329 |
+
else:
|
| 330 |
+
return clean_p
|
| 331 |
+
|
| 332 |
+
# 【修复】使用 *args 接收参数,避免 Gradio 传参顺序问题
|
| 333 |
+
def run_inference(*args):
|
| 334 |
+
global is_interrupted
|
| 335 |
+
is_interrupted = False
|
| 336 |
+
|
| 337 |
+
# 解析参数顺序
|
| 338 |
+
# [prompt, checks(N), weights(N), t, v, perf_mode, w, h, steps, cfg, seed, random, batch, vram_th]
|
| 339 |
+
idx = 0
|
| 340 |
+
prompt = args[idx]; idx += 1
|
| 341 |
+
num_loras = len(LORA_FILES)
|
| 342 |
+
lora_checks = args[idx : idx+num_loras]; idx += num_loras
|
| 343 |
+
lora_weights = args[idx : idx+num_loras]; idx += num_loras
|
| 344 |
+
|
| 345 |
+
t_choice = args[idx]; idx += 1
|
| 346 |
+
v_choice = args[idx]; idx += 1
|
| 347 |
+
perf_mode = args[idx]; idx += 1 # 【新增】性能模式
|
| 348 |
+
|
| 349 |
+
w = args[idx]; idx += 1
|
| 350 |
+
h = args[idx]; idx += 1
|
| 351 |
+
steps = args[idx]; idx += 1
|
| 352 |
+
cfg = args[idx]; idx += 1
|
| 353 |
+
seed = args[idx]; idx += 1
|
| 354 |
+
is_random = args[idx]; idx += 1
|
| 355 |
+
batch_size = args[idx]; idx += 1
|
| 356 |
+
vram_threshold = args[idx]; idx += 1
|
| 357 |
+
|
| 358 |
+
auto_flush_vram(vram_threshold)
|
| 359 |
+
clean_w = (int(w) // 16) * 16
|
| 360 |
+
clean_h = (int(h) // 16) * 16
|
| 361 |
+
|
| 362 |
+
selected_loras, weights_map = process_lora_inputs(lora_checks, lora_weights)
|
| 363 |
+
|
| 364 |
+
# 构建最终 Prompt
|
| 365 |
+
if selected_loras:
|
| 366 |
+
tags = []
|
| 367 |
+
for f in selected_loras:
|
| 368 |
+
w_val = weights_map.get(f, 1.0)
|
| 369 |
+
name = os.path.splitext(f)[0]
|
| 370 |
+
tags.append(f"<lora:{name}:{w_val:.2f}>")
|
| 371 |
+
clean_p = re.sub(r"\s*<lora:[^>]+>", "", prompt or "").strip()
|
| 372 |
+
final_prompt = f"{clean_p} {' '.join(tags)}"
|
| 373 |
+
else:
|
| 374 |
+
final_prompt = prompt
|
| 375 |
+
|
| 376 |
+
try:
|
| 377 |
+
# 【修改】传入 perf_mode
|
| 378 |
+
pipe = manager.get_pipeline(t_choice, v_choice, selected_loras, weights_map, mode='txt', perf_mode=perf_mode)
|
| 379 |
+
except Exception as e:
|
| 380 |
+
raise gr.Error(f"模型加载失败: {str(e)}")
|
| 381 |
+
|
| 382 |
+
if is_random: seed = random.randint(0, 2**32 - 1)
|
| 383 |
+
generator = torch.Generator(DEVICE).manual_seed(int(seed))
|
| 384 |
+
|
| 385 |
+
date_folder = datetime.now().strftime("%Y-%m-%d")
|
| 386 |
+
save_dir = os.path.join(OUTPUT_ROOT, date_folder)
|
| 387 |
+
os.makedirs(save_dir, exist_ok=True)
|
| 388 |
+
|
| 389 |
+
results_images = []
|
| 390 |
+
progress = gr.Progress()
|
| 391 |
+
|
| 392 |
+
try:
|
| 393 |
+
print(f"🔥 任务启动 | 图片分辨率: {clean_w}x{clean_h} | 种子: {seed}")
|
| 394 |
+
step_callback = make_progress_callback(progress, int(steps))
|
| 395 |
+
|
| 396 |
+
for i in range(int(batch_size)):
|
| 397 |
+
if is_interrupted: break
|
| 398 |
+
output = pipe(
|
| 399 |
+
prompt=final_prompt,
|
| 400 |
+
width=clean_w,
|
| 401 |
+
height=clean_h,
|
| 402 |
+
num_inference_steps=int(steps),
|
| 403 |
+
guidance_scale=float(cfg),
|
| 404 |
+
generator=generator,
|
| 405 |
+
callback_on_step_end=step_callback
|
| 406 |
+
).images[0]
|
| 407 |
+
|
| 408 |
+
filename = f"{datetime.now().strftime('%H%M%S')}_{uuid.uuid4().hex[:4]}.png"
|
| 409 |
+
path = os.path.join(save_dir, filename)
|
| 410 |
+
output.save(path)
|
| 411 |
+
results_images.append(output)
|
| 412 |
+
_, current_status = get_vram_info()
|
| 413 |
+
yield results_images, seed, current_status
|
| 414 |
+
|
| 415 |
+
except Exception as e:
|
| 416 |
+
if "任务已手动停止" in str(e):
|
| 417 |
+
print("🛑 任务已停止")
|
| 418 |
+
else:
|
| 419 |
+
import traceback
|
| 420 |
+
traceback.print_exc()
|
| 421 |
+
raise gr.Error(f"生成中断: {str(e)}")
|
| 422 |
+
finally:
|
| 423 |
+
# del pipe
|
| 424 |
+
auto_flush_vram(vram_threshold)
|
| 425 |
+
|
| 426 |
+
# 【修复】图生图
|
| 427 |
+
def run_img2img(*args, progress=gr.Progress()):
|
| 428 |
+
global is_interrupted
|
| 429 |
+
is_interrupted = False
|
| 430 |
+
|
| 431 |
+
# [input_image, prompt, checks(N), weights(N), t, v, perf_mode, ...fixed...]
|
| 432 |
+
idx = 0
|
| 433 |
+
input_image = args[idx]; idx += 1
|
| 434 |
+
prompt = args[idx]; idx += 1
|
| 435 |
+
|
| 436 |
+
num_loras = len(LORA_FILES)
|
| 437 |
+
lora_checks = args[idx : idx+num_loras]; idx += num_loras
|
| 438 |
+
lora_weights = args[idx : idx+num_loras]; idx += num_loras
|
| 439 |
+
|
| 440 |
+
t_choice = args[idx]; idx += 1
|
| 441 |
+
v_choice = args[idx]; idx += 1
|
| 442 |
+
perf_mode = args[idx]; idx += 1 # 【新增】性能模式
|
| 443 |
+
|
| 444 |
+
output_width = args[idx]; idx += 1
|
| 445 |
+
output_height = args[idx]; idx += 1
|
| 446 |
+
strength = args[idx]; idx += 1
|
| 447 |
+
steps = args[idx]; idx += 1
|
| 448 |
+
cfg = args[idx]; idx += 1
|
| 449 |
+
seed = args[idx]; idx += 1
|
| 450 |
+
is_random = args[idx]; idx += 1
|
| 451 |
+
batch_size = args[idx]; idx += 1
|
| 452 |
+
vram_threshold = args[idx]; idx += 1
|
| 453 |
+
|
| 454 |
+
if input_image is None:
|
| 455 |
+
raise gr.Error("❌ 请先上传图片")
|
| 456 |
+
|
| 457 |
+
auto_flush_vram(vram_threshold)
|
| 458 |
+
selected_loras, weights_map = process_lora_inputs(lora_checks, lora_weights)
|
| 459 |
+
|
| 460 |
+
if selected_loras:
|
| 461 |
+
tags = []
|
| 462 |
+
for f in selected_loras:
|
| 463 |
+
w_val = weights_map.get(f, 1.0)
|
| 464 |
+
name = os.path.splitext(f)[0]
|
| 465 |
+
tags.append(f"<lora:{name}:{w_val:.2f}>")
|
| 466 |
+
clean_p = re.sub(r"\s*<lora:[^>]+>", "", prompt or "").strip()
|
| 467 |
+
final_prompt = f"{clean_p} {' '.join(tags)}"
|
| 468 |
+
else:
|
| 469 |
+
final_prompt = prompt
|
| 470 |
+
|
| 471 |
+
if output_width == 0 or output_height == 0:
|
| 472 |
+
orig_w, orig_h = input_image.size
|
| 473 |
+
aspect = orig_w / orig_h
|
| 474 |
+
target_size = 1024
|
| 475 |
+
if aspect > 1:
|
| 476 |
+
target_w, target_h = target_size, max(512, int(target_size / aspect))
|
| 477 |
+
else:
|
| 478 |
+
target_h, target_w = target_size, max(512, int(target_size * aspect))
|
| 479 |
+
target_w = (target_w // 16) * 16
|
| 480 |
+
target_h = (target_h // 16) * 16
|
| 481 |
+
else:
|
| 482 |
+
target_w = (int(output_width) // 16) * 16
|
| 483 |
+
target_h = (int(output_height) // 16) * 16
|
| 484 |
+
|
| 485 |
+
input_image = input_image.convert("RGB").resize((target_w, target_h))
|
| 486 |
+
|
| 487 |
+
if is_random: seed = random.randint(0, 2**32 - 1)
|
| 488 |
+
generator = torch.Generator(DEVICE).manual_seed(int(seed))
|
| 489 |
+
|
| 490 |
+
date_folder = datetime.now().strftime("%Y-%m-%d")
|
| 491 |
+
save_dir = os.path.join(OUTPUT_ROOT, date_folder)
|
| 492 |
+
os.makedirs(save_dir, exist_ok=True)
|
| 493 |
+
|
| 494 |
+
results = []
|
| 495 |
+
pipe = None
|
| 496 |
+
|
| 497 |
+
try:
|
| 498 |
+
# 【修改】传入 perf_mode
|
| 499 |
+
pipe = manager.get_pipeline(t_choice, v_choice, selected_loras, weights_map, mode='img', perf_mode=perf_mode)
|
| 500 |
+
|
| 501 |
+
for i in progress.tqdm(range(int(batch_size)), desc="图生图生成中"):
|
| 502 |
+
if is_interrupted: break
|
| 503 |
+
torch.cuda.ipc_collect()
|
| 504 |
+
step_callback = make_progress_callback(progress, int(steps))
|
| 505 |
+
|
| 506 |
+
output = pipe(
|
| 507 |
+
prompt=final_prompt,
|
| 508 |
+
image=input_image,
|
| 509 |
+
strength=float(strength),
|
| 510 |
+
num_inference_steps=int(steps),
|
| 511 |
+
guidance_scale=0.0,
|
| 512 |
+
generator=generator,
|
| 513 |
+
callback_on_step_end=step_callback
|
| 514 |
+
).images[0]
|
| 515 |
+
|
| 516 |
+
filename = f"img2img_{datetime.now().strftime('%H%M%S')}_{uuid.uuid4().hex[:4]}.png"
|
| 517 |
+
path = os.path.join(save_dir, filename)
|
| 518 |
+
output.save(path)
|
| 519 |
+
results.append(path)
|
| 520 |
+
|
| 521 |
+
except Exception as e:
|
| 522 |
+
if "任务已手动停止" in str(e):
|
| 523 |
+
print("🛑 任务已停止")
|
| 524 |
+
else:
|
| 525 |
+
import traceback
|
| 526 |
+
traceback.print_exc()
|
| 527 |
+
raise gr.Error(f"生成中断: {str(e)}")
|
| 528 |
+
finally:
|
| 529 |
+
del pipe
|
| 530 |
+
auto_flush_vram(vram_threshold)
|
| 531 |
+
_, current_status = get_vram_info()
|
| 532 |
+
|
| 533 |
+
return results, seed, current_status
|
| 534 |
+
|
| 535 |
+
# 【修复】融合图
|
| 536 |
+
def run_fusion_img(*args, progress=gr.Progress()):
|
| 537 |
+
global is_interrupted
|
| 538 |
+
is_interrupted = False
|
| 539 |
+
|
| 540 |
+
# [image1, image2, prompt, checks(N), weights(N), t, v, perf_mode, ...fixed...]
|
| 541 |
+
idx = 0
|
| 542 |
+
image1 = args[idx]; idx += 1
|
| 543 |
+
image2 = args[idx]; idx += 1
|
| 544 |
+
prompt = args[idx]; idx += 1
|
| 545 |
+
|
| 546 |
+
num_loras = len(LORA_FILES)
|
| 547 |
+
lora_checks = args[idx : idx+num_loras]; idx += num_loras
|
| 548 |
+
lora_weights = args[idx : idx+num_loras]; idx += num_loras
|
| 549 |
+
|
| 550 |
+
t_choice = args[idx]; idx += 1
|
| 551 |
+
v_choice = args[idx]; idx += 1
|
| 552 |
+
perf_mode = args[idx]; idx += 1 # 【新增】性能模式
|
| 553 |
+
|
| 554 |
+
output_width = args[idx]; idx += 1
|
| 555 |
+
output_height = args[idx]; idx += 1
|
| 556 |
+
blend_strength = args[idx]; idx += 1
|
| 557 |
+
strength = args[idx]; idx += 1
|
| 558 |
+
steps = args[idx]; idx += 1
|
| 559 |
+
cfg = args[idx]; idx += 1
|
| 560 |
+
seed = args[idx]; idx += 1
|
| 561 |
+
is_random = args[idx]; idx += 1
|
| 562 |
+
batch_size = args[idx]; idx += 1
|
| 563 |
+
vram_threshold = args[idx]; idx += 1
|
| 564 |
+
|
| 565 |
+
if image1 is None or image2 is None:
|
| 566 |
+
raise gr.Error("❌ 请上传两张参考图片")
|
| 567 |
+
|
| 568 |
+
auto_flush_vram(vram_threshold)
|
| 569 |
+
selected_loras, weights_map = process_lora_inputs(lora_checks, lora_weights)
|
| 570 |
+
|
| 571 |
+
if selected_loras:
|
| 572 |
+
tags = []
|
| 573 |
+
for f in selected_loras:
|
| 574 |
+
w_val = weights_map.get(f, 1.0)
|
| 575 |
+
name = os.path.splitext(f)[0]
|
| 576 |
+
tags.append(f"<lora:{name}:{w_val:.2f}>")
|
| 577 |
+
clean_p = re.sub(r"\s*<lora:[^>]+>", "", prompt or "").strip()
|
| 578 |
+
final_prompt = f"{clean_p} {' '.join(tags)}"
|
| 579 |
+
else:
|
| 580 |
+
final_prompt = prompt
|
| 581 |
+
|
| 582 |
+
if output_width == 0 or output_height == 0:
|
| 583 |
+
orig_w, orig_h = image1.size
|
| 584 |
+
aspect = orig_w / orig_h
|
| 585 |
+
target_size = 1024
|
| 586 |
+
if aspect > 1:
|
| 587 |
+
target_w, target_h = target_size, max(512, int(target_size / aspect))
|
| 588 |
+
else:
|
| 589 |
+
target_h, target_w = target_size, max(512, int(target_size * aspect))
|
| 590 |
+
target_w = (target_w // 16) * 16
|
| 591 |
+
target_h = (target_h // 16) * 16
|
| 592 |
+
else:
|
| 593 |
+
target_w = (int(output_width) // 16) * 16
|
| 594 |
+
target_h = (int(output_height) // 16) * 16
|
| 595 |
+
|
| 596 |
+
image1 = image1.convert("RGB").resize((target_w, target_h))
|
| 597 |
+
image2 = image2.convert("RGB").resize((target_w, target_h))
|
| 598 |
+
blended_image = Image.blend(image1, image2, float(blend_strength))
|
| 599 |
+
|
| 600 |
+
if is_random: seed = random.randint(0, 2**32 - 1)
|
| 601 |
+
generator = torch.Generator(DEVICE).manual_seed(int(seed))
|
| 602 |
+
|
| 603 |
+
date_folder = datetime.now().strftime("%Y-%m-%d")
|
| 604 |
+
save_dir = os.path.join(OUTPUT_ROOT, date_folder)
|
| 605 |
+
os.makedirs(save_dir, exist_ok=True)
|
| 606 |
+
|
| 607 |
+
results = []
|
| 608 |
+
pipe = None
|
| 609 |
+
|
| 610 |
+
try:
|
| 611 |
+
# 【修改】传入 perf_mode
|
| 612 |
+
pipe = manager.get_pipeline(t_choice, v_choice, selected_loras, weights_map, mode='img', perf_mode=perf_mode)
|
| 613 |
+
|
| 614 |
+
for i in progress.tqdm(range(int(batch_size)), desc="融合生成中"):
|
| 615 |
+
if is_interrupted: break
|
| 616 |
+
torch.cuda.ipc_collect()
|
| 617 |
+
step_callback = make_progress_callback(progress, int(steps))
|
| 618 |
+
|
| 619 |
+
output = pipe(
|
| 620 |
+
prompt=final_prompt,
|
| 621 |
+
image=blended_image,
|
| 622 |
+
strength=float(strength),
|
| 623 |
+
num_inference_steps=int(steps),
|
| 624 |
+
guidance_scale=0.0,
|
| 625 |
+
generator=generator,
|
| 626 |
+
callback_on_step_end=step_callback
|
| 627 |
+
).images[0]
|
| 628 |
+
|
| 629 |
+
filename = f"fusion_{datetime.now().strftime('%H%M%S')}_{uuid.uuid4().hex[:4]}.png"
|
| 630 |
+
path = os.path.join(save_dir, filename)
|
| 631 |
+
output.save(path)
|
| 632 |
+
results.append(path)
|
| 633 |
+
|
| 634 |
+
except Exception as e:
|
| 635 |
+
if "任务已手动停止" in str(e):
|
| 636 |
+
print("🛑 任务已停止")
|
| 637 |
+
else:
|
| 638 |
+
import traceback
|
| 639 |
+
traceback.print_exc()
|
| 640 |
+
raise gr.Error(f"生成中断: {str(e)}")
|
| 641 |
+
finally:
|
| 642 |
+
del pipe
|
| 643 |
+
auto_flush_vram(vram_threshold)
|
| 644 |
+
_, current_status = get_vram_info()
|
| 645 |
+
|
| 646 |
+
return results, seed, current_status
|
| 647 |
+
|
| 648 |
+
# ==========================================
|
| 649 |
+
# UI 界面
|
| 650 |
+
# ==========================================
|
| 651 |
+
js_kill_window = """
|
| 652 |
+
function() {
|
| 653 |
+
setTimeout(function(){ window.close(); }, 1000);
|
| 654 |
+
document.body.innerHTML = '<div style="display:flex;justify-content:center;align-items:center;height:100vh;background:#000;color:#fff;font-family:sans-serif;"><h1>🚫 系统已关闭,请直接关闭此标签页</h1></div>';
|
| 655 |
+
document.body.style.backgroundColor = "black";
|
| 656 |
+
document.body.style.overflow = "hidden";
|
| 657 |
+
return [];
|
| 658 |
+
}
|
| 659 |
+
"""
|
| 660 |
+
|
| 661 |
+
def kill_system_process():
|
| 662 |
+
print("🛑 正在执行一键退出程序...")
|
| 663 |
+
try:
|
| 664 |
+
os.system("taskkill /F /IM Z-Image-Launcher.exe")
|
| 665 |
+
except Exception:
|
| 666 |
+
pass
|
| 667 |
+
time.sleep(1)
|
| 668 |
+
try:
|
| 669 |
+
os.system("taskkill /F /IM python.exe")
|
| 670 |
+
except Exception:
|
| 671 |
+
pass
|
| 672 |
+
sys.exit(0)
|
| 673 |
+
|
| 674 |
+
# 【新增】根据显存��动判断默认模式
|
| 675 |
+
DEFAULT_PERF_MODE = "低端机 (显存优化)" if TOTAL_VRAM < 20 * 1024**3 else "高端机 (显存>=20GB)"
|
| 676 |
+
|
| 677 |
+
with gr.Blocks(title="造相 Z-Image Pro Studio | 作者: ") as demo:
|
| 678 |
+
|
| 679 |
+
print('\n' + '!'*60)
|
| 680 |
+
print(' 本软件由 Leewheel 免费分享,严禁售卖!')
|
| 681 |
+
print('!'*60 + '\n')
|
| 682 |
+
gr.Warning('本软件由 Leewheel 免费分享。如果你是付费购买,你被骗了!', duration=20)
|
| 683 |
+
with gr.Row(elem_id="header_row"):
|
| 684 |
+
gr.Markdown("# 🎨 造相 Z-Image Pro Studio | 作者: Leewheel(V1.00C)")
|
| 685 |
+
exit_btn = gr.Button("❌ 一键退出系统", variant="stop", scale=0, min_width=150)
|
| 686 |
+
|
| 687 |
+
vram_info_display = gr.Markdown("显存状态加载中...")
|
| 688 |
+
|
| 689 |
+
with gr.Tabs():
|
| 690 |
+
# --- 文成图 ---
|
| 691 |
+
with gr.Tab("文成图"):
|
| 692 |
+
with gr.Row():
|
| 693 |
+
with gr.Column(scale=4):
|
| 694 |
+
prompt_input = gr.Textbox(label="Prompt", lines=4)
|
| 695 |
+
manual_flush_btn = gr.Button("🧹 清理显存", size="sm", variant="secondary")
|
| 696 |
+
vram_threshold_slider = gr.Slider(50, 98, 90, step=1, label="自动清理阈值 (%)")
|
| 697 |
+
|
| 698 |
+
# 【修改】LoRA 区域添加刷新按钮
|
| 699 |
+
with gr.Accordion("LoRA 权重设置 (每个 LoRA 独立调节)", open=False):
|
| 700 |
+
txt_lora_checks = []
|
| 701 |
+
txt_lora_sliders = []
|
| 702 |
+
|
| 703 |
+
# 【新增】LoRA 刷新按钮和状态提示
|
| 704 |
+
with gr.Row():
|
| 705 |
+
txt_refresh_lora_btn = gr.Button("🔄 刷新 LoRA 文件列表", size="sm")
|
| 706 |
+
txt_lora_info_md = gr.Markdown("")
|
| 707 |
+
|
| 708 |
+
if not LORA_FILES:
|
| 709 |
+
gr.Markdown("*未检测到 LoRA 文件*")
|
| 710 |
+
else:
|
| 711 |
+
for fname in LORA_FILES:
|
| 712 |
+
with gr.Row():
|
| 713 |
+
chk = gr.Checkbox(label=fname, value=False, scale=1, container=False)
|
| 714 |
+
sld = gr.Slider(0, 2.0, 1.0, step=0.05, label="权重", scale=4)
|
| 715 |
+
txt_lora_checks.append(chk)
|
| 716 |
+
txt_lora_sliders.append(sld)
|
| 717 |
+
|
| 718 |
+
with gr.Accordion("模型设置", open=True):
|
| 719 |
+
refresh_models_btn = gr.Button("🔄 刷新底模/VAE", size="sm")
|
| 720 |
+
t_drop = gr.Dropdown(label="Transformer", choices=["default"] + scan_model_items(MOD_TRANS_DIR), value="default")
|
| 721 |
+
v_drop = gr.Dropdown(label="VAE", choices=["default"] + scan_model_items(MOD_VAE_DIR), value="default")
|
| 722 |
+
|
| 723 |
+
# 【新增】性能模式选项
|
| 724 |
+
perf_mode_radio = gr.Radio(
|
| 725 |
+
choices=["高端机 (显存>=20GB)", "低端机 (显存优化)"],
|
| 726 |
+
value=DEFAULT_PERF_MODE,
|
| 727 |
+
label="性能模式"
|
| 728 |
+
)
|
| 729 |
+
|
| 730 |
+
with gr.Row():
|
| 731 |
+
width_s = gr.Slider(512, 2048, 1024, step=16, label="宽 (16倍数)")
|
| 732 |
+
height_s = gr.Slider(512, 2048, 1024, step=16, label="高 (16倍数)")
|
| 733 |
+
step_s = gr.Slider(1, 50, 8, label="步数")
|
| 734 |
+
cfg_s = gr.Slider(0, 10, 0, label="CFG")
|
| 735 |
+
batch_s = gr.Slider(1, 32, 1, step=1, label="生成张数")
|
| 736 |
+
seed_n = gr.Number(label="种子", value=42, precision=0)
|
| 737 |
+
random_c = gr.Checkbox(label="随机种子", value=True)
|
| 738 |
+
|
| 739 |
+
with gr.Row():
|
| 740 |
+
run_btn = gr.Button("🚀 开始生成", variant="primary", size="lg")
|
| 741 |
+
stop_btn = gr.Button("🛑 停止生成", variant="stop", size="lg", interactive=False)
|
| 742 |
+
|
| 743 |
+
with gr.Column(scale=6):
|
| 744 |
+
res_gallery = gr.Gallery(label="输出结果", columns=2, height="80vh")
|
| 745 |
+
res_seed = gr.Number(label="种子", interactive=False)
|
| 746 |
+
vram_info_display = gr.Markdown("显存状态加载中...")
|
| 747 |
+
|
| 748 |
+
# --- 图片编辑 ---
|
| 749 |
+
with gr.Tab("图片编辑"):
|
| 750 |
+
with gr.Row():
|
| 751 |
+
with gr.Column():
|
| 752 |
+
image_input = gr.Image(label="上传图片", type="pil")
|
| 753 |
+
with gr.Group():
|
| 754 |
+
rotate_angle = gr.Slider(-360, 360, 0, step=1, label="旋转角度 (度)")
|
| 755 |
+
crop_x = gr.Slider(0, 100, 0, step=1, label="裁剪 X (%)")
|
| 756 |
+
crop_y = gr.Slider(0, 100, 0, step=1, label="裁剪 Y (%)")
|
| 757 |
+
crop_width = gr.Slider(0, 100, 100, step=1, label="裁剪宽度 (%)")
|
| 758 |
+
crop_height = gr.Slider(0, 100, 100, step=1, label="裁剪高度 (%)")
|
| 759 |
+
flip_horizontal = gr.Checkbox(label="水平翻转")
|
| 760 |
+
flip_vertical = gr.Checkbox(label="垂直翻转")
|
| 761 |
+
edit_btn = gr.Button("开始编辑", variant="primary")
|
| 762 |
+
with gr.Column():
|
| 763 |
+
edited_image_output = gr.Image(label="编辑后的图片", type="pil")
|
| 764 |
+
with gr.Group():
|
| 765 |
+
apply_filter = gr.Dropdown(["模糊", "轮廓", "细节", "边缘增强", "更多边缘增强", "浮雕", "查找边缘", "锐化", "平滑", "更多平滑"], label="应用滤镜")
|
| 766 |
+
brightness = gr.Slider(-100, 100, 0, step=1, label="亮度调整 (%)")
|
| 767 |
+
contrast = gr.Slider(-100, 100, 0, step=1, label="对比度调整 (%)")
|
| 768 |
+
saturation = gr.Slider(-100, 100, 0, step=1, label="饱和度调整 (%)")
|
| 769 |
+
|
| 770 |
+
def edit_image(image, angle, x, y, width, height, hflip, vflip, filter, brightness, contrast, saturation):
|
| 771 |
+
if image is None: return None
|
| 772 |
+
if angle != 0: image = image.rotate(angle, expand=True)
|
| 773 |
+
if x or y or width < 100 or height < 100:
|
| 774 |
+
original_width, original_height = image.size
|
| 775 |
+
left = int(original_width * x / 100)
|
| 776 |
+
top = int(original_height * y / 100)
|
| 777 |
+
right = int(original_width * (x + width) / 100)
|
| 778 |
+
bottom = int(original_height * (y + height) / 100)
|
| 779 |
+
image = image.crop((left, top, right, bottom))
|
| 780 |
+
if hflip: image = ImageOps.mirror(image)
|
| 781 |
+
if vflip: image = ImageOps.flip(image)
|
| 782 |
+
if filter:
|
| 783 |
+
filter_map = {
|
| 784 |
+
"模糊": ImageFilter.BLUR, "轮廓": ImageFilter.CONTOUR, "细节": ImageFilter.DETAIL,
|
| 785 |
+
"边缘增强": ImageFilter.EDGE_ENHANCE, "更多边缘增强": ImageFilter.EDGE_ENHANCE_MORE,
|
| 786 |
+
"浮雕": ImageFilter.EMBOSS, "查找边缘": ImageFilter.FIND_EDGES,
|
| 787 |
+
"锐化": ImageFilter.SHARPEN, "平滑": ImageFilter.SMOOTH, "更多平滑": ImageFilter.SMOOTH_MORE
|
| 788 |
+
}
|
| 789 |
+
filter_func = filter_map.get(filter)
|
| 790 |
+
if filter_func: image = image.filter(filter_func)
|
| 791 |
+
if brightness != 0:
|
| 792 |
+
enhancer = ImageEnhance.Brightness(image)
|
| 793 |
+
image = enhancer.enhance(1 + brightness / 100)
|
| 794 |
+
if contrast != 0:
|
| 795 |
+
enhancer = ImageEnhance.Contrast(image)
|
| 796 |
+
image = enhancer.enhance(1 + contrast / 100)
|
| 797 |
+
if saturation != 0:
|
| 798 |
+
enhancer = ImageEnhance.Color(image)
|
| 799 |
+
image = enhancer.enhance(1 + saturation / 100)
|
| 800 |
+
return image
|
| 801 |
+
|
| 802 |
+
edit_btn.click(
|
| 803 |
+
fn=edit_image,
|
| 804 |
+
inputs=[image_input, rotate_angle, crop_x, crop_y, crop_width, crop_height, flip_horizontal, flip_vertical, apply_filter, brightness, contrast, saturation],
|
| 805 |
+
outputs=edited_image_output
|
| 806 |
+
)
|
| 807 |
+
|
| 808 |
+
# --- 图生图 ---
|
| 809 |
+
with gr.Tab("图生图"):
|
| 810 |
+
with gr.Row():
|
| 811 |
+
with gr.Column(scale=4):
|
| 812 |
+
with gr.Group():
|
| 813 |
+
img2img_input = gr.Image(label="上传参考图", type="pil")
|
| 814 |
+
img2img_prompt = gr.Textbox(label="Prompt (推荐)", lines=2, placeholder="描述你想要生成的画面...")
|
| 815 |
+
img2img_flush = gr.Button("🧹 清理显存", size="sm", variant="secondary")
|
| 816 |
+
|
| 817 |
+
# 【修改】图生图独立控件,添加刷新
|
| 818 |
+
with gr.Accordion("LoRA 权重设置 (独立调节)", open=False):
|
| 819 |
+
i2i_lora_checks = []
|
| 820 |
+
i2i_lora_sliders = []
|
| 821 |
+
|
| 822 |
+
with gr.Row():
|
| 823 |
+
i2i_refresh_lora_btn = gr.Button("🔄 刷新 LoRA 文件列表", size="sm")
|
| 824 |
+
i2i_lora_info_md = gr.Markdown("")
|
| 825 |
+
|
| 826 |
+
if not LORA_FILES:
|
| 827 |
+
gr.Markdown("*未检测到 LoRA 文件*")
|
| 828 |
+
else:
|
| 829 |
+
for fname in LORA_FILES:
|
| 830 |
+
with gr.Row():
|
| 831 |
+
chk = gr.Checkbox(label=fname, value=False, scale=1, container=False)
|
| 832 |
+
sld = gr.Slider(0, 2.0, 1.0, step=0.05, label="权重", scale=4)
|
| 833 |
+
i2i_lora_checks.append(chk)
|
| 834 |
+
i2i_lora_sliders.append(sld)
|
| 835 |
+
|
| 836 |
+
with gr.Accordion("模型与参数", open=True):
|
| 837 |
+
img2img_refresh_models = gr.Button("🔄 ��新底模/VAE", size="sm")
|
| 838 |
+
img2img_t_drop = gr.Dropdown(label="Transformer", choices=["default"] + scan_model_items(MOD_TRANS_DIR), value="default")
|
| 839 |
+
img2img_v_drop = gr.Dropdown(label="VAE", choices=["default"] + scan_model_items(MOD_VAE_DIR), value="default")
|
| 840 |
+
|
| 841 |
+
# 【新增】图生图性能模式
|
| 842 |
+
img2img_perf_mode = gr.Radio(
|
| 843 |
+
choices=["高端机 (显存>=20GB)", "低端机 (显存优化)"],
|
| 844 |
+
value=DEFAULT_PERF_MODE,
|
| 845 |
+
label="性能模式"
|
| 846 |
+
)
|
| 847 |
+
|
| 848 |
+
with gr.Row():
|
| 849 |
+
img2img_width_s = gr.Slider(0, 2048, 0, step=16, label="输出宽 (0=自动保持比例)")
|
| 850 |
+
img2img_height_s = gr.Slider(0, 2048, 0, step=16, label="输出高 (0=自动保持比例)")
|
| 851 |
+
gr.Markdown("**提示:** 宽高都为0时自动保持上传图比例并接近1024;手动设置大于512时生效")
|
| 852 |
+
img2img_strength = gr.Slider(0.0, 1.0, 0.75, step=0.01, label="重绘强度")
|
| 853 |
+
img2img_steps = gr.Slider(1, 100, 12, step=1, label="步数")
|
| 854 |
+
img2img_cfg = gr.Number(value=0.0, label="CFG(Turbo模型固定为0.0)", interactive=False)
|
| 855 |
+
img2img_batch = gr.Slider(1, 8, 1, step=1, label="张数")
|
| 856 |
+
img2img_seed = gr.Number(label="种子", value=42, precision=0)
|
| 857 |
+
img2img_random = gr.Checkbox(label="随机种子", value=True)
|
| 858 |
+
with gr.Row():
|
| 859 |
+
img2img_run_btn = gr.Button("🚀 生成", variant="primary", size="lg")
|
| 860 |
+
img2img_stop_btn = gr.Button("🛑 停止", variant="stop", size="lg", interactive=False)
|
| 861 |
+
with gr.Column(scale=6):
|
| 862 |
+
img2img_gallery = gr.Gallery(label="图生图结果", columns=2, height="80vh")
|
| 863 |
+
img2img_res_seed = gr.Number(label="种子", interactive=False)
|
| 864 |
+
|
| 865 |
+
# --- 融合图 ---
|
| 866 |
+
with gr.Tab("融合图"):
|
| 867 |
+
gr.Markdown("**融合2张图片**:图片1提供主要结构/姿势,图片2提供细节/脸部/风格。")
|
| 868 |
+
with gr.Row():
|
| 869 |
+
with gr.Column(scale=4):
|
| 870 |
+
with gr.Group():
|
| 871 |
+
fusion_input1 = gr.Image(label="图片1(主结构/姿势)", type="pil")
|
| 872 |
+
fusion_input2 = gr.Image(label="图片2(细节/脸部/风格)", type="pil")
|
| 873 |
+
fusion_prompt = gr.Textbox(label="融合描述 Prompt", lines=3)
|
| 874 |
+
fusion_flush = gr.Button("🧹 清理显存", size="sm", variant="secondary")
|
| 875 |
+
|
| 876 |
+
# 【修改】融合图独立控件,添加刷新
|
| 877 |
+
with gr.Accordion("LoRA 权重设置 (独立调节)", open=False):
|
| 878 |
+
fusion_lora_checks = []
|
| 879 |
+
fusion_lora_sliders = []
|
| 880 |
+
|
| 881 |
+
with gr.Row():
|
| 882 |
+
fusion_refresh_lora_btn = gr.Button("🔄 刷新 LoRA 文件列表", size="sm")
|
| 883 |
+
fusion_lora_info_md = gr.Markdown("")
|
| 884 |
+
|
| 885 |
+
if not LORA_FILES:
|
| 886 |
+
gr.Markdown("*未检测到 LoRA 文件*")
|
| 887 |
+
else:
|
| 888 |
+
for fname in LORA_FILES:
|
| 889 |
+
with gr.Row():
|
| 890 |
+
chk = gr.Checkbox(label=fname, value=False, scale=1, container=False)
|
| 891 |
+
sld = gr.Slider(0, 2.0, 1.0, step=0.05, label="权重", scale=4)
|
| 892 |
+
fusion_lora_checks.append(chk)
|
| 893 |
+
fusion_lora_sliders.append(sld)
|
| 894 |
+
|
| 895 |
+
with gr.Accordion("模型与参数", open=True):
|
| 896 |
+
fusion_refresh_models = gr.Button("🔄 刷新底模/VAE", size="sm")
|
| 897 |
+
fusion_t_drop = gr.Dropdown(label="Transformer", choices=["default"] + scan_model_items(MOD_TRANS_DIR), value="default")
|
| 898 |
+
fusion_v_drop = gr.Dropdown(label="VAE", choices=["default"] + scan_model_items(MOD_VAE_DIR), value="default")
|
| 899 |
+
|
| 900 |
+
# 【新增】融合图性能模式
|
| 901 |
+
fusion_perf_mode = gr.Radio(
|
| 902 |
+
choices=["高端机 (显存>=20GB)", "低端机 (显存优化)"],
|
| 903 |
+
value=DEFAULT_PERF_MODE,
|
| 904 |
+
label="性能模式"
|
| 905 |
+
)
|
| 906 |
+
|
| 907 |
+
with gr.Row():
|
| 908 |
+
fusion_width_s = gr.Slider(0, 2048, 0, step=16, label="输出宽 (0=自动保持比例)")
|
| 909 |
+
fusion_height_s = gr.Slider(0, 2048, 0, step=16, label="输出高 (0=自动保持比例)")
|
| 910 |
+
gr.Markdown("**提示:** 宽高都为0时自动保持图片1比例并接近1024")
|
| 911 |
+
with gr.Row():
|
| 912 |
+
fusion_blend = gr.Slider(0.0, 1.0, 0.5, step=0.05, label="图片2混合强度 (0=全用图片1, 1=全用图片2)")
|
| 913 |
+
fusion_strength = gr.Slider(0.0, 1.0, 0.7, step=0.05, label="重绘强度 (越高变化越大)")
|
| 914 |
+
fusion_steps = gr.Slider(1, 100, 15, step=1, label="步数")
|
| 915 |
+
fusion_cfg = gr.Number(value=0.0, label="CFG(固定为0.0)", interactive=False)
|
| 916 |
+
fusion_batch = gr.Slider(1, 8, 1, step=1, label="张数")
|
| 917 |
+
fusion_seed = gr.Number(label="种子", value=42, precision=0)
|
| 918 |
+
fusion_random = gr.Checkbox(label="随机种子", value=True)
|
| 919 |
+
with gr.Row():
|
| 920 |
+
fusion_run_btn = gr.Button("🚀 开始融合", variant="primary", size="lg")
|
| 921 |
+
fusion_stop_btn = gr.Button("🛑 停止", variant="stop", size="lg", interactive=False)
|
| 922 |
+
with gr.Column(scale=6):
|
| 923 |
+
fusion_gallery = gr.Gallery(label="融合结果", columns=2, height="80vh")
|
| 924 |
+
fusion_res_seed = gr.Number(label="种子", interactive=False)
|
| 925 |
+
|
| 926 |
+
# -----------------------
|
| 927 |
+
# UI状态函数
|
| 928 |
+
# -----------------------
|
| 929 |
+
def ui_to_running():
|
| 930 |
+
return gr.update(interactive=False), gr.update(interactive=True)
|
| 931 |
+
|
| 932 |
+
def ui_to_idle():
|
| 933 |
+
return gr.update(interactive=True), gr.update(interactive=False)
|
| 934 |
+
|
| 935 |
+
def trigger_interrupt():
|
| 936 |
+
global is_interrupted
|
| 937 |
+
is_interrupted = True
|
| 938 |
+
return "🛑 正在强制中断..."
|
| 939 |
+
|
| 940 |
+
# -----------------------
|
| 941 |
+
# 按钮事件绑定
|
| 942 |
+
# -----------------------
|
| 943 |
+
|
| 944 |
+
# 退出按钮
|
| 945 |
+
exit_btn.click(fn=kill_system_process, js=js_kill_window)
|
| 946 |
+
|
| 947 |
+
# 文生图
|
| 948 |
+
refresh_models_btn.click(
|
| 949 |
+
fn=lambda: (
|
| 950 |
+
gr.update(choices=["default"] + scan_model_items(MOD_TRANS_DIR)),
|
| 951 |
+
gr.update(choices=["default"] + scan_model_items(MOD_VAE_DIR))
|
| 952 |
+
),
|
| 953 |
+
outputs=[t_drop, v_drop]
|
| 954 |
+
)
|
| 955 |
+
|
| 956 |
+
# 【新增】文生图 LoRA 刷新
|
| 957 |
+
txt_refresh_lora_btn.click(fn=refresh_lora_list, outputs=txt_lora_info_md)
|
| 958 |
+
|
| 959 |
+
manual_flush_btn.click(
|
| 960 |
+
fn=lambda: (gc.collect(), torch.cuda.empty_cache(), get_vram_info()[1])[2],
|
| 961 |
+
outputs=vram_info_display
|
| 962 |
+
)
|
| 963 |
+
|
| 964 |
+
txt_ui_inputs = [prompt_input] + txt_lora_checks + txt_lora_sliders
|
| 965 |
+
for c in txt_lora_checks + txt_lora_sliders:
|
| 966 |
+
c.change(fn=update_prompt_ui_base, inputs=txt_ui_inputs, outputs=prompt_input)
|
| 967 |
+
|
| 968 |
+
inference_event = run_btn.click(
|
| 969 |
+
fn=ui_to_running,
|
| 970 |
+
outputs=[run_btn, stop_btn]
|
| 971 |
+
).then(
|
| 972 |
+
fn=run_inference,
|
| 973 |
+
# 【修改】输入增加 perf_mode_radio
|
| 974 |
+
inputs=txt_ui_inputs + [t_drop, v_drop, perf_mode_radio, width_s, height_s, step_s, cfg_s, seed_n, random_c, batch_s, vram_threshold_slider],
|
| 975 |
+
outputs=[res_gallery, res_seed, vram_info_display]
|
| 976 |
+
).then(
|
| 977 |
+
fn=ui_to_idle,
|
| 978 |
+
outputs=[run_btn, stop_btn]
|
| 979 |
+
)
|
| 980 |
+
|
| 981 |
+
stop_btn.click(
|
| 982 |
+
fn=trigger_interrupt,
|
| 983 |
+
outputs=vram_info_display
|
| 984 |
+
).then(
|
| 985 |
+
fn=ui_to_idle,
|
| 986 |
+
outputs=[run_btn, stop_btn],
|
| 987 |
+
cancels=[inference_event]
|
| 988 |
+
)
|
| 989 |
+
|
| 990 |
+
# 图生图
|
| 991 |
+
def refresh_all_models_img():
|
| 992 |
+
return gr.update(choices=["default"] + scan_model_items(MOD_TRANS_DIR)), gr.update(choices=["default"] + scan_model_items(MOD_VAE_DIR))
|
| 993 |
+
img2img_refresh_models.click(fn=refresh_all_models_img, outputs=[img2img_t_drop, img2img_v_drop])
|
| 994 |
+
|
| 995 |
+
# 【新增】图生图 LoRA 刷新
|
| 996 |
+
i2i_refresh_lora_btn.click(fn=refresh_lora_list, outputs=i2i_lora_info_md)
|
| 997 |
+
|
| 998 |
+
img2img_flush.click(fn=lambda: (gc.collect(), torch.cuda.empty_cache(), get_vram_info()[1])[2], outputs=vram_info_display)
|
| 999 |
+
|
| 1000 |
+
i2i_ui_inputs = [img2img_prompt] + i2i_lora_checks + i2i_lora_sliders
|
| 1001 |
+
for c in i2i_lora_checks + i2i_lora_sliders:
|
| 1002 |
+
c.change(fn=update_prompt_ui_base, inputs=i2i_ui_inputs, outputs=img2img_prompt)
|
| 1003 |
+
|
| 1004 |
+
img2img_event = img2img_run_btn.click(fn=ui_to_running, outputs=[img2img_run_btn, img2img_stop_btn])\
|
| 1005 |
+
.then(fn=run_img2img,
|
| 1006 |
+
# 【修改】输入增加 img2img_perf_mode
|
| 1007 |
+
inputs=[img2img_input, img2img_prompt] + i2i_lora_checks + i2i_lora_sliders +
|
| 1008 |
+
[img2img_t_drop, img2img_v_drop, img2img_perf_mode, img2img_width_s, img2img_height_s,
|
| 1009 |
+
img2img_strength, img2img_steps, img2img_cfg, img2img_seed, img2img_random, img2img_batch, vram_threshold_slider],
|
| 1010 |
+
outputs=[img2img_gallery, img2img_res_seed, vram_info_display])\
|
| 1011 |
+
.then(fn=ui_to_idle, outputs=[img2img_run_btn, img2img_stop_btn])
|
| 1012 |
+
|
| 1013 |
+
img2img_stop_btn.click(fn=trigger_interrupt, outputs=vram_info_display).then(fn=ui_to_idle, outputs=[img2img_run_btn, img2img_stop_btn], cancels=[img2img_event])
|
| 1014 |
+
|
| 1015 |
+
# 融合图
|
| 1016 |
+
fusion_refresh_models.click(fn=refresh_all_models_img, outputs=[fusion_t_drop, fusion_v_drop])
|
| 1017 |
+
|
| 1018 |
+
# 【新增】融合图 LoRA 刷新
|
| 1019 |
+
fusion_refresh_lora_btn.click(fn=refresh_lora_list, outputs=fusion_lora_info_md)
|
| 1020 |
+
|
| 1021 |
+
fusion_flush.click(fn=lambda: (gc.collect(), torch.cuda.empty_cache(), get_vram_info()[1])[2], outputs=vram_info_display)
|
| 1022 |
+
|
| 1023 |
+
fusion_ui_inputs = [fusion_prompt] + fusion_lora_checks + fusion_lora_sliders
|
| 1024 |
+
for c in fusion_lora_checks + fusion_lora_sliders:
|
| 1025 |
+
c.change(fn=update_prompt_ui_base, inputs=fusion_ui_inputs, outputs=fusion_prompt)
|
| 1026 |
+
|
| 1027 |
+
fusion_event = fusion_run_btn.click(fn=ui_to_running, outputs=[fusion_run_btn, fusion_stop_btn])\
|
| 1028 |
+
.then(fn=run_fusion_img,
|
| 1029 |
+
# 【修改】输入增加 fusion_perf_mode
|
| 1030 |
+
inputs=[fusion_input1, fusion_input2, fusion_prompt] + fusion_lora_checks + fusion_lora_sliders +
|
| 1031 |
+
[fusion_t_drop, fusion_v_drop, fusion_perf_mode, fusion_width_s, fusion_height_s,
|
| 1032 |
+
fusion_blend, fusion_strength, fusion_steps, fusion_cfg,
|
| 1033 |
+
fusion_seed, fusion_random, fusion_batch, vram_threshold_slider],
|
| 1034 |
+
outputs=[fusion_gallery, fusion_res_seed, vram_info_display])\
|
| 1035 |
+
.then(fn=ui_to_idle, outputs=[fusion_run_btn, fusion_stop_btn])
|
| 1036 |
+
|
| 1037 |
+
fusion_stop_btn.click(fn=trigger_interrupt, outputs=vram_info_display).then(fn=ui_to_idle, outputs=[fusion_run_btn, fusion_stop_btn], cancels=[fusion_event])
|
| 1038 |
+
|
| 1039 |
+
if __name__ == "__main__":
|
| 1040 |
+
demo.launch(share=False, inbrowser=True)
|
assets/DMDR.webp
ADDED
|
Git LFS Details
|
assets/Z-Image-Gallery.pdf
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:6f9895b3246d2547bac74bbe0be975da500eaae93f2cad4248ad3281786b1ac6
|
| 3 |
+
size 15767436
|
assets/architecture.webp
ADDED
|
Git LFS Details
|
assets/decoupled-dmd.webp
ADDED
|
Git LFS Details
|
assets/image_arena_all.jpg
ADDED
|
Git LFS Details
|
assets/image_arena_os.jpg
ADDED
|
Git LFS Details
|
assets/leaderboard.png
ADDED
|
Git LFS Details
|
assets/reasoning.png
ADDED
|
Git LFS Details
|
assets/showcase.jpg
ADDED
|
Git LFS Details
|
assets/showcase_editing.png
ADDED
|
Git LFS Details
|
assets/showcase_realistic.png
ADDED
|
assets/showcase_rendering.png
ADDED
|
Git LFS Details
|
assets/training_pipeline.jpg
ADDED
|
Git LFS Details
|