Update imagic.py
Browse files
imagic.py
CHANGED
@@ -291,6 +291,8 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
|
|
291 |
text_embeddings.requires_grad_(False)
|
292 |
|
293 |
# Now we fine tune the unet to better reconstruct the image
|
|
|
|
|
294 |
self.unet.requires_grad_(True)
|
295 |
self.unet.train()
|
296 |
optimizer = torch.optim.Adam(
|
|
|
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(
|