修复加载bug

README.md

@@ -52,7 +52,7 @@ Place the LoRA file under `lora/` first (or set `LORA_PATH`); otherwise the app
 - Prompt
 - Resolution category + explicit WxH selection
 - Seed (with random toggle)
-- Steps, time shift, max sequence length
+- Steps, CFG, scheduler + shift (and extra scheduler params), max sequence length
 - LoRA toggle + strength (enabled only if the file is found)
 
 ## Git LFS note

app.py

@@ -3,6 +3,7 @@ import random
 import re
 import threading
 import warnings
+import inspect
 from typing import List, Tuple
 
 import gradio as gr
@@ -22,6 +23,7 @@ OFFLOAD_TO_CPU_AFTER_RUN = os.environ.get("OFFLOAD_TO_CPU_AFTER_RUN", "true").lo
 ENABLE_AOTI = os.environ.get("ENABLE_AOTI", "false").lower() == "true"
 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"))
 
 warnings.filterwarnings("ignore")
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -101,6 +103,14 @@ pipe_lock = threading.Lock()
 pipe_on_gpu: bool = False
 aoti_loaded: bool = False
 
+SCHEDULERS = {"FlowMatch Euler": FlowMatchEulerDiscreteScheduler}
+try:
+    from diffusers import FlowMatchHeunDiscreteScheduler  # type: ignore
+
+    SCHEDULERS["FlowMatch Heun"] = FlowMatchHeunDiscreteScheduler
+except Exception:
+    pass
+
 
 def parse_resolution(resolution: str) -> Tuple[int, int]:
     match = re.search(r"(\d+)\s*[×x]\s*(\d+)", resolution)
@@ -109,6 +119,30 @@ def parse_resolution(resolution: str) -> Tuple[int, int]:
     return 1024, 1024
 
 
+def set_attention_backend_safe(transformer, backend: str) -> str:
+    candidates: List[str] = []
+    if backend:
+        candidates.append(backend)
+        if backend.startswith("_"):
+            candidates.append(backend.lstrip("_"))
+        else:
+            candidates.append(f"_{backend}")
+    candidates.extend(["flash", "xformers", "native"])
+
+    last_exc: Exception | None = None
+    for name in candidates:
+        if not name:
+            continue
+        try:
+            transformer.set_attention_backend(name)
+            return name
+        except Exception as exc:  # noqa: BLE001
+            last_exc = exc
+            continue
+
+    raise RuntimeError(f"Failed to set attention backend (tried {candidates}): {last_exc}")
+
+
 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"
@@ -132,7 +166,7 @@ def set_lora_scale(pipeline: ZImagePipeline, scale: float) -> None:
 
 def load_models() -> Tuple[ZImagePipeline, bool, str | None]:
     global pipe, lora_loaded, lora_error
-    if pipe is not None:
+    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
@@ -163,7 +197,7 @@ def load_models() -> Tuple[ZImagePipeline, bool, str | None]:
 
     tokenizer.padding_side = "left"
 
-    pipe = ZImagePipeline(scheduler=None, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=None)
+    pipeline = ZImagePipeline(scheduler=None, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=None)
 
     if not os.path.exists(MODEL_PATH):
         transformer = ZImageTransformer2DModel.from_pretrained(
@@ -178,26 +212,31 @@ def load_models() -> Tuple[ZImagePipeline, bool, str | None]:
             torch_dtype=torch.bfloat16,
         )
 
-    transformer.set_attention_backend(ATTENTION_BACKEND)
+    applied_backend = set_attention_backend_safe(transformer, ATTENTION_BACKEND)
+    print(f"Attention backend: {applied_backend}")
 
-    pipe.transformer = transformer
+    pipeline.transformer = transformer
 
-    lora_loaded, lora_error = attach_lora(pipe)
+    loaded, error = attach_lora(pipeline)
+    lora_loaded, lora_error = loaded, error
     if lora_error:
         print(lora_error)
     else:
         print(f"LoRA loaded: {lora_loaded} ({LORA_PATH})")
 
+    pipe = pipeline
     return pipe, lora_loaded, lora_error
 
 
 def ensure_models_loaded() -> Tuple[ZImagePipeline, bool, str | None]:
-    global pipe
-    if pipe is not None:
+    global pipe, pipe_on_gpu
+    if pipe is not None and getattr(pipe, "transformer", None) is not None:
         return pipe, lora_loaded, lora_error
     with pipe_lock:
-        if pipe is not None:
+        if pipe is not None and getattr(pipe, "transformer", None) is not None:
             return pipe, lora_loaded, lora_error
+        pipe = None
+        pipe_on_gpu = False
         return load_models()
 
 
@@ -205,6 +244,8 @@ def ensure_on_gpu() -> None:
     global pipe_on_gpu, aoti_loaded
     if pipe is None:
         raise gr.Error("Model not loaded.")
+    if getattr(pipe, "transformer", None) is None:
+        raise gr.Error("Model init failed (transformer missing). Check startup logs.")
     if not torch.cuda.is_available():
         raise gr.Error("CUDA is not available. This Space requires a GPU.")
     if pipe_on_gpu:
@@ -241,8 +282,33 @@ def offload_to_cpu() -> None:
         torch.cuda.empty_cache()
 
 
-def set_scheduler(pipeline: ZImagePipeline, shift: float) -> None:
-    scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift)
+def make_scheduler(scheduler_cls, **kwargs):
+    sig = inspect.signature(scheduler_cls.__init__)
+    accepted = set(sig.parameters.keys())
+    accepted.discard("self")
+    filtered = {k: v for k, v in kwargs.items() if k in accepted and v is not None}
+    return scheduler_cls(**filtered)
+
+
+def set_scheduler(
+    pipeline: ZImagePipeline,
+    scheduler_name: str,
+    *,
+    num_train_timesteps: int,
+    shift: float,
+    use_dynamic_shifting: bool,
+    base_shift: float,
+    max_shift: float,
+) -> None:
+    scheduler_cls = SCHEDULERS.get(scheduler_name, FlowMatchEulerDiscreteScheduler)
+    scheduler = make_scheduler(
+        scheduler_cls,
+        num_train_timesteps=int(num_train_timesteps),
+        shift=float(shift),
+        use_dynamic_shifting=bool(use_dynamic_shifting),
+        base_shift=float(base_shift),
+        max_shift=float(max_shift),
+    )
     pipeline.scheduler = scheduler
 
 
@@ -257,10 +323,23 @@ def generate_image(
     max_sequence_length: int,
     use_lora: bool,
     lora_scale: float,
+    scheduler_name: str,
+    num_train_timesteps: int,
+    use_dynamic_shifting: bool,
+    base_shift: float,
+    max_shift: float,
 ) -> Tuple[torch.Tensor, int]:
     width, height = parse_resolution(resolution)
     generator = torch.Generator("cuda").manual_seed(seed)
-    set_scheduler(pipeline, shift)
+    set_scheduler(
+        pipeline,
+        scheduler_name,
+        num_train_timesteps=num_train_timesteps,
+        shift=shift,
+        use_dynamic_shifting=use_dynamic_shifting,
+        base_shift=base_shift,
+        max_shift=max_shift,
+    )
 
     if lora_loaded:
         if use_lora:
@@ -327,6 +406,12 @@ def generate(
     seed: int = 42,
     steps: int = 9,
     shift: float = 3.0,
+    cfg: float = DEFAULT_CFG,
+    scheduler_name: str = "FlowMatch Euler",
+    num_train_timesteps: int = 1000,
+    use_dynamic_shifting: bool = False,
+    base_shift: float = 0.5,
+    max_shift: float = 3.0,
     random_seed: bool = True,
     use_lora: bool = True,
     lora_scale: float = 1.0,
@@ -347,10 +432,15 @@ def generate(
             seed=new_seed,
             steps=int(steps),
             shift=float(shift),
-            guidance_scale=0.0,
+            guidance_scale=float(cfg),
             max_sequence_length=int(max_sequence_length),
             use_lora=use_lora,
             lora_scale=float(lora_scale),
+            scheduler_name=str(scheduler_name),
+            num_train_timesteps=int(num_train_timesteps),
+            use_dynamic_shifting=bool(use_dynamic_shifting),
+            base_shift=float(base_shift),
+            max_shift=float(max_shift),
         )[0]
     finally:
         if OFFLOAD_TO_CPU_AFTER_RUN:
@@ -397,11 +487,33 @@ Model: `{MODEL_PATH}` | {pipe_status}
                 seed = gr.Number(label="Seed", value=42, precision=0)
                 random_seed = gr.Checkbox(label="Random Seed", value=True)
 
-            with gr.Row():
-                steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=9, step=1)
-                shift = gr.Slider(label="Time Shift", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
+            with gr.Accordion("KSampler / Advanced", open=False):
+                with gr.Row():
+                    steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=9, step=1)
+                    cfg = gr.Slider(label="CFG", minimum=0.0, maximum=10.0, value=DEFAULT_CFG, step=0.1)
+
+                with gr.Row():
+                    scheduler_name = gr.Dropdown(
+                        label="Scheduler",
+                        choices=list(SCHEDULERS.keys()),
+                        value="FlowMatch Euler",
+                    )
+                    num_train_timesteps = gr.Slider(
+                        label="num_train_timesteps",
+                        minimum=100,
+                        maximum=2000,
+                        value=1000,
+                        step=10,
+                    )
+
+                with gr.Row():
+                    shift = gr.Slider(label="Shift", minimum=0.0, maximum=10.0, value=3.0, step=0.1)
+                    use_dynamic_shifting = gr.Checkbox(label="use_dynamic_shifting", value=False)
+
+                with gr.Row():
+                    base_shift = gr.Slider(label="base_shift", minimum=0.0, maximum=10.0, value=0.5, step=0.1)
+                    max_shift = gr.Slider(label="max_shift", minimum=0.0, maximum=10.0, value=3.0, step=0.1)
 
-            with gr.Row():
                 max_seq = gr.Slider(label="Max Sequence Length", minimum=256, maximum=1024, value=512, step=16)
 
             with gr.Row():
@@ -443,7 +555,24 @@ Model: `{MODEL_PATH}` | {pipe_status}
 
     generate_btn.click(
         generate,
-        inputs=[prompt_input, resolution, seed, steps, shift, random_seed, use_lora, lora_strength, max_seq, output_gallery],
+        inputs=[
+            prompt_input,
+            resolution,
+            seed,
+            steps,
+            shift,
+            cfg,
+            scheduler_name,
+            num_train_timesteps,
+            use_dynamic_shifting,
+            base_shift,
+            max_shift,
+            random_seed,
+            use_lora,
+            lora_strength,
+            max_seq,
+            output_gallery,
+        ],
         outputs=[output_gallery, used_seed, seed],
         api_visibility="public",
     )