tolgacangoz
commited on
Commit
•
a5673c8
1
Parent(s):
987a636
Upload matryoshka.py
Browse files- scheduler/matryoshka.py +17 -5
scheduler/matryoshka.py
CHANGED
@@ -660,12 +660,24 @@ class MatryoshkaDDIMScheduler(SchedulerMixin, ConfigMixin):
|
|
660 |
)
|
661 |
|
662 |
if variance_noise is None:
|
663 |
-
|
664 |
-
|
665 |
-
|
666 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
667 |
|
668 |
-
|
669 |
|
670 |
if not return_dict:
|
671 |
return (prev_sample,)
|
|
|
660 |
)
|
661 |
|
662 |
if variance_noise is None:
|
663 |
+
if len(model_output) > 1:
|
664 |
+
variance_noise = []
|
665 |
+
for m_o in model_output:
|
666 |
+
variance_noise.append(
|
667 |
+
randn_tensor(
|
668 |
+
m_o.shape, generator=generator, device=m_o.device, dtype=m_o.dtype
|
669 |
+
)
|
670 |
+
)
|
671 |
+
else:
|
672 |
+
variance_noise = randn_tensor(
|
673 |
+
model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype
|
674 |
+
)
|
675 |
+
if len(model_output) > 1:
|
676 |
+
prev_sample = [p_s + std_dev_t * v_n for v_n, p_s in zip(variance_noise, prev_sample)]
|
677 |
+
else:
|
678 |
+
variance = std_dev_t * variance_noise
|
679 |
|
680 |
+
prev_sample = prev_sample + variance
|
681 |
|
682 |
if not return_dict:
|
683 |
return (prev_sample,)
|