Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 from diffusers import UniPCMultistepScheduler
 from diffusers import WanPipeline, AutoencoderKLWan
+from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
 from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
 from huggingface_hub import hf_hub_download
 from PIL import Image
@@ -13,7 +14,20 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # --- MODEL SETUP ---
 model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
 vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
-pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
+transformer=WanTransformer3DModel.from_pretrained(model_id,
+        subfolder='transformer',
+        torch_dtype=torch.bfloat16,
+        device_map='cuda',
+    )
+transformer_2=WanTransformer3DModel.from_pretrained(model_id,
+        subfolder='transformer_2',
+        torch_dtype=torch.bfloat16,
+        device_map='cuda',
+    )
+pipe = WanPipeline.from_pretrained(model_id,
+                                   transformer=transformer,
+                                   transformer_2 = transformer_2,
+                                   vae=vae, torch_dtype=torch.bfloat16)
 flow_shift = 2.0
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
 pipe.to(device)
@@ -146,7 +160,8 @@ def generate(prompt, negative_prompt, width, height, num_inference_steps, option
         num_frames=1,
         num_inference_steps=num_inference_steps,
         guidance_scale=float(guidance_scale),
-        # guidance_scale_2=float(guidance_scale_2),
+        guidance_scale_2=float(guidance_scale_2),
+        boundary_ratio=0.3,
     )
     image = output.frames[0][0]
     image = (image * 255).astype(np.uint8)