Update app.py

app.py

@@ -10,9 +10,10 @@ from transformers import AutoTokenizer, Qwen3ForCausalLM
 from controlnet_aux.processor import Processor
 from PIL import Image
 from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
 
 # Import pipeline and model
-# Ensure videox_fun is in your python path
+# Ensure the videox_fun folder is in your current directory
 from videox_fun.pipeline import ZImageControlPipeline
 from videox_fun.models import ZImageControlTransformer2DModel
 
@@ -24,23 +25,24 @@ except ImportError:
     def polish_prompt(prompt):
         return prompt
 
-# Configuration
+# --- Configuration & Paths ---
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1280
 
-# Paths
-MODEL_LOCAL = "models/Z-Image-Turbo/" # Local path or HuggingFace ID
-# We prioritize the local safetensors file for ControlNet weights
-CONTROLNET_WEIGHTS = "models/Z-Image-Turbo-Fun-Controlnet-Union.safetensors"
+# Hugging Face Repo IDs
+MODEL_REPO = "Tongyi-MAI/Z-Image-Turbo"
+CONTROLNET_REPO = "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union"
+CONTROLNET_FILENAME = "Z-Image-Turbo-Fun-Controlnet-Union.safetensors"
 
-print("Loading Z-Image Turbo model...")
+print(f"Loading Z-Image Turbo from {MODEL_REPO}...")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 weight_dtype = torch.bfloat16
 
-# 1. Load Transformer with Control Config
+# --- 1. Load Transformer ---
 print("Initializing Transformer...")
+# We load the config and model structure from the main Hub Repo
 transformer = ZImageControlTransformer2DModel.from_pretrained(
-    MODEL_LOCAL,
+    MODEL_REPO,
     subfolder="transformer",
     transformer_additional_kwargs={
         "control_layers_places": [0, 5, 10, 15, 20, 25],
@@ -48,8 +50,22 @@ transformer = ZImageControlTransformer2DModel.from_pretrained(
     },
 ).to(device, weight_dtype)
 
-# 2. Load ControlNet Weights manually
-if os.path.exists(CONTROLNET_WEIGHTS):
+# --- 2. Download & Load ControlNet Weights ---
+# Check if weights exist locally; if not, download them
+if not os.path.exists(CONTROLNET_FILENAME):
+    print(f"Downloading ControlNet weights from {CONTROLNET_REPO}...")
+    try:
+        CONTROLNET_WEIGHTS = hf_hub_download(
+            repo_id=CONTROLNET_REPO,
+            filename=CONTROLNET_FILENAME
+        )
+    except Exception as e:
+        print(f"Failed to download ControlNet weights: {e}")
+        CONTROLNET_WEIGHTS = None
+else:
+    CONTROLNET_WEIGHTS = CONTROLNET_FILENAME
+
+if CONTROLNET_WEIGHTS:
     print(f"Loading ControlNet weights from {CONTROLNET_WEIGHTS}")
     try:
         state_dict = load_file(CONTROLNET_WEIGHTS)
@@ -61,32 +77,32 @@ if os.path.exists(CONTROLNET_WEIGHTS):
     except Exception as e:
         print(f"Error loading ControlNet weights: {e}")
 else:
-    print(f"Warning: ControlNet weights not found at {CONTROLNET_WEIGHTS}. Trying to run without them or using base weights.")
+    print("Warning: Running without explicit ControlNet weights.")
 
-# 3. Load VAE, Tokenizer, Encoder, Scheduler
-print("Loading core components...")
+# --- 3. Load Core Components ---
+print("Loading VAE, Tokenizer, and Text Encoder...")
 vae = AutoencoderKL.from_pretrained(
-    MODEL_LOCAL,
+    MODEL_REPO,
     subfolder="vae",
 ).to(device, weight_dtype)
 
 tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_LOCAL, 
+    MODEL_REPO, 
     subfolder="tokenizer"
 )
 
 text_encoder = Qwen3ForCausalLM.from_pretrained(
-    MODEL_LOCAL, 
+    MODEL_REPO, 
     subfolder="text_encoder", 
     torch_dtype=weight_dtype,
 ).to(device)
 
 scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-    MODEL_LOCAL, 
+    MODEL_REPO, 
     subfolder="scheduler"
 )
 
-# 4. Assemble Pipeline
+# --- 4. Assemble Pipeline ---
 pipe = ZImageControlPipeline(
     vae=vae,
     tokenizer=tokenizer,
@@ -284,7 +300,7 @@ button.primary:hover {
 }
 """
 
-with gr.Blocks(title="Z-Image Turbo ControlNet") as demo:
+with gr.Blocks(title="Z-Image Turbo ControlNet", css=apple_css) as demo:
     
     gr.HTML("""
         <div class="header-container">
@@ -369,5 +385,4 @@ with gr.Blocks(title="Z-Image Turbo ControlNet") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(share=False,
-               css=apple_css)
+    demo.launch(share=False)