22

app.py

@@ -27,13 +27,13 @@ weight_dtype = torch.bfloat16
 transformer = ZImageControlTransformer2DModel.from_pretrained(
     MODEL_LOCAL,
     subfolder="transformer",
-    low_cpu_mem_usage=False,
-    dtype=torch.bfloat16,
+    # low_cpu_mem_usage=False,
+    # torch_dtype=torch.bfloat16,
     transformer_additional_kwargs={
         "control_layers_places": [0, 5, 10, 15, 20, 25],
         "control_in_dim": 16
     },
-).to(torch.bfloat16)
+).to("cuda", torch.bfloat16)
 
 if TRANSFORMER_LOCAL is not None:
     print(f"From checkpoint: {TRANSFORMER_LOCAL}")
@@ -69,8 +69,9 @@ pipe = ZImageControlPipeline(
     transformer=transformer,
     scheduler=scheduler,
 )
-pipe.transformer = transformer
-pipe.to("cuda")
+pipe.to("cuda", torch.bfloat16)
+# pipe.transformer = transformer
+# pipe.to("cuda")
 
 # ======== AoTI compilation + FA3 ========
 pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]