修复PERT和AIOT

.gitignore

@@ -0,0 +1 @@
+AGENTS.md

README.md

@@ -26,16 +26,18 @@ Gradio Space using the official Z-Image pipeline (`Tongyi-MAI/Z-Image-Turbo`) wi
 3) Manually add the LoRA file from https://civitai.com/models/2206377/zit-mystic-xxx to `lora/zit-mystic-xxx.safetensors` (or set `LORA_PATH`). Network fetch of Civitai is not handled in the Space.  
 4) If model download fails with a token error, set `HF_TOKEN` in the Space secrets (some repos require authentication).  
 5) (Optional) Toggle advanced envs below; then the Space will launch `app.py`. The header shows whether the LoRA was detected/loaded.
+   - If the header/log says `PEFT backend is required for LoRA`, install `peft` (already included in `requirements.txt`) and restart/rebuild.
 
 ## Environment variables
 - `MODEL_PATH` (default `Tongyi-MAI/Z-Image-Turbo`): HF repo or local path for the Z-Image model.
 - `LORA_PATH` (default `lora/zit-mystic-xxx.safetensors`): Path to the LoRA file; loaded if present.
 - `HF_TOKEN`: HF token for gated/private models or faster pulls.
+- `MODEL_DTYPE` (default `auto`): `bf16` if supported, else `fp16` (override with `bf16`/`fp16`/`fp32`).
 - `ENABLE_COMPILE` (default `false`): Enable `torch.compile` on the transformer.
 - `ENABLE_WARMUP` (default `false`): Run a quick warmup across resolutions after load (adds startup time).
-- `ATTENTION_BACKEND` (default `flash_3`): Backend for transformer attention.
-- `OFFLOAD_TO_CPU_AFTER_RUN` (default `true`): Move the model back to CPU after each generation to play nicer with ZeroGPU.
-- `ENABLE_AOTI` (default `false`): Try to load ZeroGPU AoTI blocks via `spaces.aoti_blocks_load` for faster inference.
+- `ATTENTION_BACKEND` (default `_flash_3`): Backend for transformer attention (falls back to `flash`/`xformers`/`native`).
+- `OFFLOAD_TO_CPU_AFTER_RUN` (default `false`): Move the model back to CPU after each generation (useful on ZeroGPU; slower on normal GPUs).
+- `ENABLE_AOTI` (default `true`): Try to load ZeroGPU AoTI blocks via `spaces.aoti_blocks_load` for faster inference.
 - `AOTI_REPO` (default `zerogpu-aoti/Z-Image`): AoTI blocks repo.
 - `AOTI_VARIANT` (default `fa3`): AoTI variant.
 

app.py

@@ -1,4 +1,5 @@
 import os
+import importlib.util
 import random
 import re
 import threading
@@ -25,6 +26,52 @@ AOTI_REPO = os.environ.get("AOTI_REPO", "zerogpu-aoti/Z-Image")
 AOTI_VARIANT = os.environ.get("AOTI_VARIANT", "fa3")
 DEFAULT_CFG = float(os.environ.get("DEFAULT_CFG", "0.0"))
 
+
+def resolve_model_dtype() -> torch.dtype:
+    override = os.environ.get("MODEL_DTYPE")
+    if override:
+        key = override.strip().lower()
+        if key in {"bf16", "bfloat16"}:
+            return torch.bfloat16
+        if key in {"fp16", "float16", "half"}:
+            return torch.float16
+        if key in {"fp32", "float32"}:
+            return torch.float32
+        print(f"Unknown MODEL_DTYPE={override!r}; falling back to auto.")
+
+    if torch.cuda.is_available():
+        is_bf16_supported = getattr(torch.cuda, "is_bf16_supported", None)
+        if callable(is_bf16_supported) and is_bf16_supported():
+            return torch.bfloat16
+        return torch.float16
+    return torch.float32
+
+
+def dtype_label(dtype: torch.dtype) -> str:
+    if dtype == torch.bfloat16:
+        return "bf16"
+    if dtype == torch.float16:
+        return "fp16"
+    if dtype == torch.float32:
+        return "fp32"
+    return str(dtype).replace("torch.", "")
+
+
+def get_gpu_summary() -> str:
+    if not torch.cuda.is_available():
+        return "CPU"
+    try:
+        name = torch.cuda.get_device_name(0)
+        major, minor = torch.cuda.get_device_capability(0)
+        return f"{name} (cc {major}.{minor})"
+    except Exception:
+        return "CUDA"
+
+
+MODEL_DTYPE = resolve_model_dtype()
+MODEL_DTYPE_LABEL = dtype_label(MODEL_DTYPE)
+GPU_SUMMARY = get_gpu_summary()
+
 if torch.cuda.is_available():
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.set_float32_matmul_precision("high")
@@ -106,6 +153,8 @@ lora_error: str | None = None
 pipe_lock = threading.Lock()
 pipe_on_gpu: bool = False
 aoti_loaded: bool = False
+applied_attention_backend: str | None = None
+aoti_error: str | None = None
 
 SCHEDULERS = {"FlowMatch Euler": FlowMatchEulerDiscreteScheduler}
 try:
@@ -116,6 +165,13 @@ except Exception:
     pass
 
 
+def module_available(module_name: str) -> bool:
+    try:
+        return importlib.util.find_spec(module_name) is not None
+    except (ImportError, ValueError):
+        return False
+
+
 def parse_resolution(resolution: str) -> Tuple[int, int]:
     match = re.search(r"(\d+)\s*[×x]\s*(\d+)", resolution)
     if match:
@@ -150,6 +206,8 @@ def set_attention_backend_safe(transformer, backend: str) -> str:
 def attach_lora(pipeline: ZImagePipeline) -> Tuple[bool, str | None]:
     if not LORA_PATH or not os.path.isfile(LORA_PATH):
         return False, "LoRA file not found"
+    if not module_available("peft"):
+        return False, "PEFT backend is required for LoRA. Install `peft` and restart."
     try:
         folder, weight_name = os.path.split(LORA_PATH)
         folder = folder or "."
@@ -169,13 +227,14 @@ def set_lora_scale(pipeline: ZImagePipeline, scale: float) -> None:
 
 
 def load_models() -> Tuple[ZImagePipeline, bool, str | None]:
-    global pipe, lora_loaded, lora_error, pipe_on_gpu
+    global pipe, lora_loaded, lora_error, pipe_on_gpu, applied_attention_backend
     if pipe is not None and getattr(pipe, "transformer", None) is not None:
         return pipe, lora_loaded, lora_error
 
     use_auth_token = HF_TOKEN if HF_TOKEN else None
     hf_kwargs = {"use_auth_token": use_auth_token} if use_auth_token else {}
     print(f"Loading Z-Image from {MODEL_PATH}...")
+    print(f"GPU: {GPU_SUMMARY} | dtype: {MODEL_DTYPE_LABEL}")
 
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is not available. This app requires a GPU.")
@@ -184,24 +243,24 @@ def load_models() -> Tuple[ZImagePipeline, bool, str | None]:
         vae = AutoencoderKL.from_pretrained(
             MODEL_PATH,
             subfolder="vae",
-            torch_dtype=torch.bfloat16,
+            torch_dtype=MODEL_DTYPE,
             **hf_kwargs,
-        ).to("cuda", torch.bfloat16)
+        ).to("cuda", MODEL_DTYPE)
         text_encoder = AutoModelForCausalLM.from_pretrained(
             MODEL_PATH,
             subfolder="text_encoder",
-            torch_dtype=torch.bfloat16,
+            torch_dtype=MODEL_DTYPE,
             **hf_kwargs,
-        ).to("cuda", torch.bfloat16).eval()
+        ).to("cuda", MODEL_DTYPE).eval()
         tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, subfolder="tokenizer", **hf_kwargs)
     else:
-        vae = AutoencoderKL.from_pretrained(os.path.join(MODEL_PATH, "vae"), torch_dtype=torch.bfloat16).to(
-            "cuda", torch.bfloat16
+        vae = AutoencoderKL.from_pretrained(os.path.join(MODEL_PATH, "vae"), torch_dtype=MODEL_DTYPE).to(
+            "cuda", MODEL_DTYPE
         )
         text_encoder = AutoModelForCausalLM.from_pretrained(
             os.path.join(MODEL_PATH, "text_encoder"),
-            torch_dtype=torch.bfloat16,
-        ).to("cuda", torch.bfloat16).eval()
+            torch_dtype=MODEL_DTYPE,
+        ).to("cuda", MODEL_DTYPE).eval()
         tokenizer = AutoTokenizer.from_pretrained(os.path.join(MODEL_PATH, "tokenizer"))
 
     tokenizer.padding_side = "left"
@@ -212,20 +271,20 @@ def load_models() -> Tuple[ZImagePipeline, bool, str | None]:
         transformer = ZImageTransformer2DModel.from_pretrained(
             MODEL_PATH,
             subfolder="transformer",
-            torch_dtype=torch.bfloat16,
+            torch_dtype=MODEL_DTYPE,
             **hf_kwargs,
         )
     else:
         transformer = ZImageTransformer2DModel.from_pretrained(
             os.path.join(MODEL_PATH, "transformer"),
-            torch_dtype=torch.bfloat16,
+            torch_dtype=MODEL_DTYPE,
         )
 
-    applied_backend = set_attention_backend_safe(transformer, ATTENTION_BACKEND)
-    print(f"Attention backend: {applied_backend}")
+    applied_attention_backend = set_attention_backend_safe(transformer, ATTENTION_BACKEND)
+    print(f"Attention backend: {applied_attention_backend}")
 
-    pipeline.transformer = transformer.to("cuda", torch.bfloat16)
-    pipeline.to("cuda", torch.bfloat16)
+    pipeline.transformer = transformer.to("cuda", MODEL_DTYPE)
+    pipeline.to("cuda", MODEL_DTYPE)
 
     loaded, error = attach_lora(pipeline)
     lora_loaded, lora_error = loaded, error
@@ -387,7 +446,7 @@ def warmup_model(pipeline: ZImagePipeline, resolutions: List[str]) -> None:
 
 
 def init_app() -> None:
-    global aoti_loaded
+    global aoti_loaded, aoti_error
     try:
         ensure_models_loaded()
         if ENABLE_AOTI and not aoti_loaded and pipe is not None and getattr(pipe, "transformer", None) is not None:
@@ -395,8 +454,10 @@ def init_app() -> None:
                 pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
                 spaces.aoti_blocks_load(pipe.transformer.layers, AOTI_REPO, variant=AOTI_VARIANT)
                 aoti_loaded = True
+                aoti_error = None
                 print(f"AoTI loaded: {AOTI_REPO} (variant={AOTI_VARIANT})")
             except Exception as exc:  # noqa: BLE001
+                aoti_error = str(exc)
                 print(f"AoTI load failed (continuing without AoTI): {exc}")
         if ENABLE_WARMUP and pipe is not None:
             ensure_on_gpu()
@@ -468,9 +529,18 @@ def generate(
 init_app()
 
 with gr.Blocks(title="Z-Image + LoRA") as demo:
-    pipe_status = "loaded (CPU)" if pipe else "not loaded"
+    pipe_status = "loaded (GPU)" if pipe and pipe_on_gpu else "loaded (CPU)" if pipe else "not loaded"
     lora_file_status = "found" if os.path.isfile(LORA_PATH) else "missing"
-    lora_status = f"LoRA file: {LORA_PATH} ({lora_file_status})"
+    if lora_loaded:
+        lora_status = f"LoRA: loaded ({LORA_PATH})"
+    elif lora_error:
+        lora_status = f"LoRA: not loaded ({lora_error})"
+    else:
+        lora_status = f"LoRA file: {LORA_PATH} ({lora_file_status})"
+
+    attention_status = applied_attention_backend or "unknown"
+    aoti_status = "loaded" if aoti_loaded else f"failed ({aoti_error})" if aoti_error else "not loaded"
+    compile_status = "on" if ENABLE_COMPILE else "off"
 
     gr.Markdown(
         f"""<div align="center">
@@ -478,6 +548,8 @@ with gr.Blocks(title="Z-Image + LoRA") as demo:
 # Z-Image Generation (No SD fallback)
 
 Model: `{MODEL_PATH}` | {pipe_status}  
+GPU: `{GPU_SUMMARY}` | dtype: `{MODEL_DTYPE_LABEL}`  
+Attention: `{attention_status}` | AoTI: `{aoti_status}` | torch.compile: `{compile_status}`  
 {lora_status}
 
 </div>"""
@@ -530,14 +602,15 @@ Model: `{MODEL_PATH}` | {pipe_status}
                 max_seq = gr.Slider(label="Max Sequence Length", minimum=256, maximum=1024, value=512, step=16)
 
             with gr.Row():
-                use_lora = gr.Checkbox(label="Use LoRA", value=True, interactive=True)
+                lora_controls_enabled = bool(lora_loaded)
+                use_lora = gr.Checkbox(label="Use LoRA", value=lora_controls_enabled, interactive=lora_controls_enabled)
                 lora_strength = gr.Slider(
                     label="LoRA Strength",
                     minimum=0.0,
                     maximum=1.5,
                     value=1.0,
                     step=0.05,
-                    interactive=True,
+                    interactive=lora_controls_enabled,
                 )
 
             generate_btn = gr.Button("Generate", variant="primary")

requirements.txt

@@ -1,6 +1,7 @@
 accelerate>=0.30.0
 diffusers>=0.32.0
 gradio>=4.44.0
+peft>=0.10.0
 Pillow>=10.0.0
 safetensors>=0.4.2
 spaces>=0.27.0