Spaces:
Running
on
Zero
Running
on
Zero
Update src/flux/xflux_pipeline.py
Browse files
src/flux/xflux_pipeline.py
CHANGED
@@ -292,8 +292,7 @@ class XFluxPipeline:
|
|
292 |
x = get_noise(
|
293 |
1, height, width, device=self.device,
|
294 |
dtype=torch.bfloat16, seed=seed
|
295 |
-
)
|
296 |
-
print("x shape:", x.shape)
|
297 |
timesteps = get_schedule(
|
298 |
num_steps,
|
299 |
(width // 8) * (height // 8) // (16 * 16),
|
@@ -312,7 +311,6 @@ class XFluxPipeline:
|
|
312 |
source_image = self.ae.encode(source_image.to(self.device).to(torch.float32))
|
313 |
# print("ae source image shape:", source_image.shape)
|
314 |
source_image = rearrange(source_image, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2).to(inp_cond['img'].dtype)
|
315 |
-
print("rearrange ae source image shape:", source_image.shape)
|
316 |
# print("rearrange ae source image shape:", source_image.shape)
|
317 |
|
318 |
if self.offload:
|
|
|
292 |
x = get_noise(
|
293 |
1, height, width, device=self.device,
|
294 |
dtype=torch.bfloat16, seed=seed
|
295 |
+
)
|
|
|
296 |
timesteps = get_schedule(
|
297 |
num_steps,
|
298 |
(width // 8) * (height // 8) // (16 * 16),
|
|
|
311 |
source_image = self.ae.encode(source_image.to(self.device).to(torch.float32))
|
312 |
# print("ae source image shape:", source_image.shape)
|
313 |
source_image = rearrange(source_image, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2).to(inp_cond['img'].dtype)
|
|
|
314 |
# print("rearrange ae source image shape:", source_image.shape)
|
315 |
|
316 |
if self.offload:
|