Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from diffusers import UniPCMultistepScheduler
+from diffusers import UniPCMultistepScheduler, FlowMatchEulerDiscreteScheduler, DDIMScheduler
 from diffusers import WanPipeline, AutoencoderKLWan  # Use Wan-specific VAE
 # from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig
 from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
@@ -19,13 +19,23 @@ flow_shift = 1.0 #5.0  1.0 for image, 5.0 for 720P, 3.0 for 480P
 # pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
 
 # Configure DDIMScheduler with a beta schedule
-pipe.scheduler = DDIMScheduler.from_config(
+# pipe.scheduler = DDIMScheduler.from_config(
+#     pipe.scheduler.config,
+#     beta_start=0.00085,  # Starting beta value
+#     beta_end=0.012,      # Ending beta value
+#     beta_schedule="linear",  # Linear beta schedule (other options: "scaled_linear", "squaredcos_cap_v2")
+#     num_train_timesteps=1000,  # Number of timesteps
+#     flow_shift=flow_shift
+# )
+
+
+# Configure FlowMatchEulerDiscreteScheduler
+pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
     pipe.scheduler.config,
-    beta_start=0.00085,  # Starting beta value
-    beta_end=0.012,      # Ending beta value
-    beta_schedule="linear",  # Linear beta schedule (other options: "scaled_linear", "squaredcos_cap_v2")
-    num_train_timesteps=1000,  # Number of timesteps
-    flow_shift=flow_shift
+    num_train_timesteps=1000,  # Standard timestep count
+    sigma_min=0.002,          # Minimum noise level
+    sigma_max=700.0,          # Maximum noise level
+    flow_shift=1.0            # Retain flow_shift for WanPipeline compatibility
 )