Update imagic.py
Browse files
imagic.py
CHANGED
@@ -272,8 +272,6 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
|
|
272 |
loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
|
273 |
accelerator.backward(loss)
|
274 |
|
275 |
-
with torch.no_grad():
|
276 |
-
torch.cuda.empty_cache()
|
277 |
optimizer.step()
|
278 |
optimizer.zero_grad()
|
279 |
|
@@ -291,8 +289,6 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
|
|
291 |
text_embeddings.requires_grad_(False)
|
292 |
|
293 |
# Now we fine tune the unet to better reconstruct the image
|
294 |
-
with torch.no_grad():
|
295 |
-
torch.cuda.empty_cache()
|
296 |
self.unet.requires_grad_(True)
|
297 |
self.unet.train()
|
298 |
optimizer = torch.optim.Adam(
|
@@ -317,12 +313,8 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
|
|
317 |
|
318 |
loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
|
319 |
accelerator.backward(loss)
|
320 |
-
|
321 |
-
|
322 |
-
torch.cuda.empty_cache()
|
323 |
-
|
324 |
-
with torch.no_grad():
|
325 |
-
optimizer.step()
|
326 |
optimizer.zero_grad()
|
327 |
|
328 |
# Checks if the accelerator has performed an optimization step behind the scenes
|
|
|
272 |
loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
|
273 |
accelerator.backward(loss)
|
274 |
|
|
|
|
|
275 |
optimizer.step()
|
276 |
optimizer.zero_grad()
|
277 |
|
|
|
289 |
text_embeddings.requires_grad_(False)
|
290 |
|
291 |
# Now we fine tune the unet to better reconstruct the image
|
|
|
|
|
292 |
self.unet.requires_grad_(True)
|
293 |
self.unet.train()
|
294 |
optimizer = torch.optim.Adam(
|
|
|
313 |
|
314 |
loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
|
315 |
accelerator.backward(loss)
|
316 |
+
|
317 |
+
optimizer.step()
|
|
|
|
|
|
|
|
|
318 |
optimizer.zero_grad()
|
319 |
|
320 |
# Checks if the accelerator has performed an optimization step behind the scenes
|