File tree Expand file tree Collapse file tree 1 file changed +4
-4
lines changed Expand file tree Collapse file tree 1 file changed +4
-4
lines changed Original file line number Diff line number Diff line change @@ -157,19 +157,19 @@ def compute_dream_and_update_latents(
157
157
with torch .no_grad ():
158
158
pred = unet (noisy_latents , timesteps , encoder_hidden_states ).sample
159
159
160
- noisy_latents , target = (None , None )
160
+ _noisy_latents , _target = (None , None )
161
161
if noise_scheduler .config .prediction_type == "epsilon" :
162
162
predicted_noise = pred
163
163
delta_noise = (noise - predicted_noise ).detach ()
164
164
delta_noise .mul_ (dream_lambda )
165
- noisy_latents = noisy_latents .add (sqrt_one_minus_alphas_cumprod * delta_noise )
166
- target = target .add (delta_noise )
165
+ _noisy_latents = noisy_latents .add (sqrt_one_minus_alphas_cumprod * delta_noise )
166
+ _target = target .add (delta_noise )
167
167
elif noise_scheduler .config .prediction_type == "v_prediction" :
168
168
raise NotImplementedError ("DREAM has not been implemented for v-prediction" )
169
169
else :
170
170
raise ValueError (f"Unknown prediction type { noise_scheduler .config .prediction_type } " )
171
171
172
- return noisy_latents , target
172
+ return _noisy_latents , _target
173
173
174
174
175
175
def unet_lora_state_dict (unet : UNet2DConditionModel ) -> Dict [str , torch .Tensor ]:
You can’t perform that action at this time.
0 commit comments