Rename policy.u() to policy.pi() to better align with the paper notation

lakonlab/models/diffusions/piflow_policies/base.py

@@ -4,7 +4,7 @@ from abc import ABCMeta, abstractmethod
 class BasePolicy(metaclass=ABCMeta):
 
     @abstractmethod
-    def u(self, x_t, sigma_t):
+    def pi(self, x_t, sigma_t):
         """Compute the flow velocity at (x_t, t).
 
         Args:

lakonlab/models/diffusions/piflow_policies/dx.py

@@ -69,7 +69,7 @@ class DXPolicy(BasePolicy):
         x_interp = (t1 - t) * x0x1[:, 0] + (t - t0) * x0x1[:, 1]
         return x_interp
 
-    def u(self, x_t, sigma_t):
+    def pi(self, x_t, sigma_t):
         """Compute the flow velocity at (x_t, t).
 
         Args:

lakonlab/models/diffusions/piflow_policies/gmflow.py

@@ -93,7 +93,7 @@ class GMFlowPolicy(BasePolicy):
             gm_vars=gm_vars,
             logweights=denoising_output['logweights'])
 
-    def u(self, x_t, sigma_t):
+    def pi(self, x_t, sigma_t):
         """Compute the flow velocity at (x_t, t).
 
         Args:

lakonlab/pipelines/piflux_pipeline.py

@@ -447,7 +447,7 @@ class PiFluxPipeline(FluxPipeline, PiFlowLoaderMixin):
                 for _ in range(num_inference_substeps[i]):
                     t = timesteps[timestep_id]
                     sigma_t = t / self.scheduler.config.num_train_timesteps
-                    u = policy.u(latents, sigma_t)
+                    u = policy.pi(latents, sigma_t)
                     latents = self.scheduler.step(u, t, latents, return_dict=False)[0]
                     timestep_id += 1
 

lakonlab/pipelines/piqwen_pipeline.py

@@ -377,7 +377,7 @@ class PiQwenImagePipeline(QwenImagePipeline, PiFlowLoaderMixin):
                 for _ in range(num_inference_substeps[i]):
                     t = timesteps[timestep_id]
                     sigma_t = t / self.scheduler.config.num_train_timesteps
-                    u = policy.u(latents, sigma_t)
+                    u = policy.pi(latents, sigma_t)
                     latents = self.scheduler.step(u, t, latents, return_dict=False)[0]
                     timestep_id += 1