From 332243f9ccc7122373e250eaf876c30fc0a06f19 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Mon, 28 Apr 2025 10:20:38 +0800
Subject: [PATCH] Update pipeline.py

---
 UNO/uno/flux/pipeline.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/UNO/uno/flux/pipeline.py b/UNO/uno/flux/pipeline.py
index 6f971f4..c144e6d 100644
--- a/UNO/uno/flux/pipeline.py
+++ b/UNO/uno/flux/pipeline.py
@@ -232,11 +232,13 @@ class UNOPipeline:
         guidance: float,
         num_steps: int,
         inp_cond={},
+        seed: int = -1,
     ):
-        # x = get_noise(
-        #     1, height, width, device=self.device,
-        #     dtype=torch.bfloat16, seed=seed
-        # )
+        x = get_noise(
+            1, height, width, device=self.device,
+            dtype=torch.bfloat16, seed=seed
+        )
+        img= rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
         timesteps = get_schedule(
             num_steps,
             (width // 8) * (height // 8) // (16 * 16),
@@ -268,6 +270,7 @@ class UNOPipeline:
 
         x = denoise(
             self.model,
+            img,
             **inp_cond,
             timesteps=timesteps,
             guidance=guidance,