stevhliu
/

one_step_unet

Model card Files Files and versions

patrickvonplaten commited on Apr 9, 2023

Commit

81b0e34

·

1 Parent(s): 3a2b2c9

Update pipeline.py

Files changed (1) hide show

pipeline.py +1 -1

pipeline.py CHANGED Viewed

@@ -11,7 +11,7 @@ class UnetSchedulerOneForwardPipeline(DiffusionPipeline):
     def __call__(self):
         image = torch.randn(
-            (1, self.unet.in_channels, self.unet.sample_size, self.unet.sample_size),
         )
         timestep = 1

     def __call__(self):
         image = torch.randn(
+            (1, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size),
         )
         timestep = 1