22

app.py

@@ -45,15 +45,18 @@ if TRANSFORMER_LOCAL is not None:
 # load ZImageControlPipeline
 vae = AutoencoderKL.from_pretrained(
     MODEL_LOCAL,
-    subfolder="vae"
+    subfolder="vae",
+    device_map="cuda"
 ).to(weight_dtype)
 
 tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_LOCAL, subfolder="tokenizer"
+    MODEL_LOCAL, 
+    subfolder="tokenizer"
 )
 text_encoder = Qwen3ForCausalLM.from_pretrained(
-    MODEL_LOCAL, subfolder="text_encoder", torch_dtype=weight_dtype,
-    low_cpu_mem_usage=False,
+    MODEL_LOCAL, 
+    subfolder="text_encoder", 
+    torch_dtype=weight_dtype,
 )
 # scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3)
 scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(