Update models/gen_pipeline.py
Browse files- models/gen_pipeline.py +1 -1
models/gen_pipeline.py
CHANGED
|
@@ -387,7 +387,7 @@ class NextStepPipeline:
|
|
| 387 |
)
|
| 388 |
|
| 389 |
# 7. unpatchify
|
| 390 |
-
latents = self.model.unpatchify(tokens)
|
| 391 |
latents = (latents / self.scaling_factor) + self.shift_factor
|
| 392 |
|
| 393 |
# 8. decode latents
|
|
|
|
| 387 |
)
|
| 388 |
|
| 389 |
# 7. unpatchify
|
| 390 |
+
latents = self.model.unpatchify(tokens, h=hw[0] // self.down_factor, w=hw[1] // self.down_factor)
|
| 391 |
latents = (latents / self.scaling_factor) + self.shift_factor
|
| 392 |
|
| 393 |
# 8. decode latents
|