尝试加速推理

app.py

@@ -18,13 +18,17 @@ LORA_PATH = os.environ.get("LORA_PATH", os.path.join("lora", "zit-mystic-xxx.saf
 HF_TOKEN = os.environ.get("HF_TOKEN")
 ENABLE_COMPILE = os.environ.get("ENABLE_COMPILE", "false").lower() == "true"
 ENABLE_WARMUP = os.environ.get("ENABLE_WARMUP", "false").lower() == "true"
-ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "flash_3")
-OFFLOAD_TO_CPU_AFTER_RUN = os.environ.get("OFFLOAD_TO_CPU_AFTER_RUN", "true").lower() == "true"
-ENABLE_AOTI = os.environ.get("ENABLE_AOTI", "false").lower() == "true"
+ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "_flash_3")
+OFFLOAD_TO_CPU_AFTER_RUN = os.environ.get("OFFLOAD_TO_CPU_AFTER_RUN", "false").lower() == "true"
+ENABLE_AOTI = os.environ.get("ENABLE_AOTI", "true").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"))
 
+if torch.cuda.is_available():
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.set_float32_matmul_precision("high")
+
 warnings.filterwarnings("ignore")
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -165,7 +169,7 @@ def set_lora_scale(pipeline: ZImagePipeline, scale: float) -> None:
 
 
 def load_models() -> Tuple[ZImagePipeline, bool, str | None]:
-    global pipe, lora_loaded, lora_error
+    global pipe, lora_loaded, lora_error, pipe_on_gpu
     if pipe is not None and getattr(pipe, "transformer", None) is not None:
         return pipe, lora_loaded, lora_error
 
@@ -173,26 +177,31 @@ def load_models() -> Tuple[ZImagePipeline, bool, str | None]:
     hf_kwargs = {"use_auth_token": use_auth_token} if use_auth_token else {}
     print(f"Loading Z-Image from {MODEL_PATH}...")
 
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is not available. This app requires a GPU.")
+
     if not os.path.exists(MODEL_PATH):
         vae = AutoencoderKL.from_pretrained(
             MODEL_PATH,
             subfolder="vae",
             torch_dtype=torch.bfloat16,
             **hf_kwargs,
-        )
+        ).to("cuda", torch.bfloat16)
         text_encoder = AutoModelForCausalLM.from_pretrained(
             MODEL_PATH,
             subfolder="text_encoder",
             torch_dtype=torch.bfloat16,
             **hf_kwargs,
-        ).eval()
+        ).to("cuda", torch.bfloat16).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)
+        vae = AutoencoderKL.from_pretrained(os.path.join(MODEL_PATH, "vae"), torch_dtype=torch.bfloat16).to(
+            "cuda", torch.bfloat16
+        )
         text_encoder = AutoModelForCausalLM.from_pretrained(
             os.path.join(MODEL_PATH, "text_encoder"),
             torch_dtype=torch.bfloat16,
-        ).eval()
+        ).to("cuda", torch.bfloat16).eval()
         tokenizer = AutoTokenizer.from_pretrained(os.path.join(MODEL_PATH, "tokenizer"))
 
     tokenizer.padding_side = "left"
@@ -215,7 +224,8 @@ def load_models() -> Tuple[ZImagePipeline, bool, str | None]:
     applied_backend = set_attention_backend_safe(transformer, ATTENTION_BACKEND)
     print(f"Attention backend: {applied_backend}")
 
-    pipeline.transformer = transformer
+    pipeline.transformer = transformer.to("cuda", torch.bfloat16)
+    pipeline.to("cuda", torch.bfloat16)
 
     loaded, error = attach_lora(pipeline)
     lora_loaded, lora_error = loaded, error
@@ -225,6 +235,7 @@ def load_models() -> Tuple[ZImagePipeline, bool, str | None]:
         print(f"LoRA loaded: {lora_loaded} ({LORA_PATH})")
 
     pipe = pipeline
+    pipe_on_gpu = True
     return pipe, lora_loaded, lora_error
 
 
@@ -241,7 +252,7 @@ def ensure_models_loaded() -> Tuple[ZImagePipeline, bool, str | None]:
 
 
 def ensure_on_gpu() -> None:
-    global pipe_on_gpu, aoti_loaded
+    global pipe_on_gpu
     if pipe is None:
         raise gr.Error("Model not loaded.")
     if getattr(pipe, "transformer", None) is None:
@@ -250,24 +261,12 @@ def ensure_on_gpu() -> None:
         raise gr.Error("CUDA is not available. This Space requires a GPU.")
     if pipe_on_gpu:
         return
-
-    print("Moving model to GPU...")
-    pipe.to("cuda", torch.bfloat16)
     pipe_on_gpu = True
 
     if ENABLE_COMPILE:
         print("Compiling transformer (torch.compile)...")
         pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs", fullgraph=False)
 
-    if ENABLE_AOTI and not aoti_loaded:
-        try:
-            pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-            spaces.aoti_blocks_load(pipe.transformer.layers, AOTI_REPO, variant=AOTI_VARIANT)
-            aoti_loaded = True
-            print(f"AoTI loaded: {AOTI_REPO} (variant={AOTI_VARIANT})")
-        except Exception as exc:  # noqa: BLE001
-            print(f"AoTI load failed (continuing without AoTI): {exc}")
-
 
 def offload_to_cpu() -> None:
     global pipe_on_gpu
@@ -388,8 +387,17 @@ def warmup_model(pipeline: ZImagePipeline, resolutions: List[str]) -> None:
 
 
 def init_app() -> None:
+    global aoti_loaded
     try:
         ensure_models_loaded()
+        if ENABLE_AOTI and not aoti_loaded and pipe is not None and getattr(pipe, "transformer", None) is not None:
+            try:
+                pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
+                spaces.aoti_blocks_load(pipe.transformer.layers, AOTI_REPO, variant=AOTI_VARIANT)
+                aoti_loaded = True
+                print(f"AoTI loaded: {AOTI_REPO} (variant={AOTI_VARIANT})")
+            except Exception as exc:  # noqa: BLE001
+                print(f"AoTI load failed (continuing without AoTI): {exc}")
         if ENABLE_WARMUP and pipe is not None:
             ensure_on_gpu()
             try: