Update app_14B.py

app_14B.py

@@ -29,19 +29,28 @@ MIN_FRAMES_MODEL = 8
 MAX_FRAMES_MODEL = 81
 
 
-pipe =  WanPipeline.from_pretrained(MODEL_ID,
-    transformer=WanTransformer3DModel.from_pretrained('linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
-        subfolder='transformer',
-        torch_dtype=torch.bfloat16,
-        device_map='cuda',
-    ),
-    transformer_2=WanTransformer3DModel.from_pretrained('linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
-        subfolder='transformer_2',
-        torch_dtype=torch.bfloat16,
-        device_map='cuda',
-    ),
-    torch_dtype=torch.bfloat16,
-).to('cuda')
+# pipe =  WanPipeline.from_pretrained(MODEL_ID,
+#     transformer=WanTransformer3DModel.from_pretrained('linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
+#         subfolder='transformer',
+#         torch_dtype=torch.bfloat16,
+#         device_map='cuda',
+#     ),
+#     transformer_2=WanTransformer3DModel.from_pretrained('linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
+#         subfolder='transformer_2',
+#         torch_dtype=torch.bfloat16,
+#         device_map='cuda',
+#     ),
+#     torch_dtype=torch.bfloat16,
+# ).to('cuda')
+
+
+vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
+pipe = WanPipeline.from_pretrained(
+    MODEL_ID, vae=vae, torch_dtype=torch.bfloat16
+)
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
+pipe.to("cuda")
+
 
 
 optimize_pipeline_(pipe,