Update app.py

app.py

@@ -7,6 +7,7 @@ import sys
 import logging
 import warnings
 import re
+import spaces
 from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
 from transformers import AutoModel, AutoTokenizer
 from dataclasses import dataclass
@@ -17,14 +18,14 @@ from diffusers import ZImagePipeline
 from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel
 from pe import prompt_template
 
-
-# ==================== Environment Variables ================================
+# ==================== Environment Variables ==================================
 MODEL_PATH = os.environ.get("MODEL_PATH", "Tongyi-MAI/Z-Image-Turbo")
 ENABLE_COMPILE = os.environ.get("ENABLE_COMPILE", "true").lower() == "true"
 ENABLE_WARMUP = os.environ.get("ENABLE_WARMUP", "true").lower() == "true"
-ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "_flash_3")
+ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "flash_3")
 DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY")
-# ===========================================================================
+HF_TOKEN = os.environ.get("HF_TOKEN")
+# =============================================================================
 
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -69,22 +70,43 @@ def get_resolution(resolution):
 
 def load_models(model_path, enable_compile=False, attention_backend="native"):
     print(f"Loading models from {model_path}...")
-    if not os.path.exists(model_path):
-        raise FileNotFoundError(f"Model directory not found: {model_path}")
-
-    vae = AutoencoderKL.from_pretrained(
-        os.path.join(model_path, "vae"), 
-        torch_dtype=torch.bfloat16, 
-        device_map="cuda"
-    )
     
-    text_encoder = AutoModel.from_pretrained(
-        os.path.join(model_path, "text_encoder"),
-        torch_dtype=torch.bfloat16,
-        device_map="cuda",
-    ).eval()
+    use_auth_token = HF_TOKEN if HF_TOKEN else True
+    
+    if not os.path.exists(model_path):
+        vae = AutoencoderKL.from_pretrained(
+            f"{model_path}/vae",
+            torch_dtype=torch.bfloat16,
+            device_map="cuda",
+            use_auth_token=use_auth_token
+        )
+        
+        text_encoder = AutoModel.from_pretrained(
+            f"{model_path}/text_encoder",
+            torch_dtype=torch.bfloat16,
+            device_map="cuda",
+            use_auth_token=use_auth_token
+        ).eval()
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            f"{model_path}/tokenizer",
+            use_auth_token=use_auth_token
+        )
+    else:
+        vae = AutoencoderKL.from_pretrained(
+            os.path.join(model_path, "vae"), 
+            torch_dtype=torch.bfloat16, 
+            device_map="cuda"
+        )
+        
+        text_encoder = AutoModel.from_pretrained(
+            os.path.join(model_path, "text_encoder"),
+            torch_dtype=torch.bfloat16,
+            device_map="cuda",
+        ).eval()
+        
+        tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"))
     
-    tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"))
     tokenizer.padding_side = "left"
     
     if enable_compile:
@@ -108,9 +130,15 @@ def load_models(model_path, enable_compile=False, attention_backend="native"):
     if enable_compile:
         pipe.vae.disable_tiling()
     
-    transformer = ZImageTransformer2DModel.from_pretrained(
-        os.path.join(model_path, "transformer")
-    ).to("cuda", torch.bfloat16)
+    if not os.path.exists(model_path):
+        transformer = ZImageTransformer2DModel.from_pretrained(
+            f"{model_path}/transformer",
+            use_auth_token=use_auth_token
+        ).to("cuda", torch.bfloat16)
+    else:
+        transformer = ZImageTransformer2DModel.from_pretrained(
+            os.path.join(model_path, "transformer")
+        ).to("cuda", torch.bfloat16)
     
     pipe.transformer = transformer
     pipe.transformer.set_attention_backend(attention_backend)
@@ -320,6 +348,7 @@ def prompt_enhance(prompt, enable_enhance):
     except Exception as e:
         return prompt, f"Error: {str(e)}"
 
+@spaces.GPU
 def generate(prompt, resolution, seed, steps, shift, enhance):
     if pipe is None:
         raise gr.Error("Model not loaded.")
@@ -350,7 +379,6 @@ def generate(prompt, resolution, seed, steps, shift, enhance):
     
     return image, final_prompt, str(seed)
 
-# ==================== Gradio Interface ====================
 init_app()
 
 with gr.Blocks(title="Z-Image Demo") as demo: