Update src/pipeline.py

src/pipeline.py

@@ -81,7 +81,7 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
             param.requires_grad = True
 
     optimizer = torch.optim.Adam([param for layer in trainable_layers for param in layer.parameters()], lr=0.001)
-
+    print(len([param for layer in trainable_layers for param in layer.parameters()]))
     instance.train()
             
     checkpoint_path= "../chkpt"
@@ -90,14 +90,15 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
         if not hasattr(pipe.vae, "_decode_without_taesd_preview"):
             pipe.vae._decode_without_taesd_preview = pipe.vae.decode
         def decode_latents_and_compare_speeds(latents, *args, **kwargs):
-            res_sd = pipe.vae._decode_without_taesd_preview(latents, *args, **kwargs)[0].cpu()
+            res_sd = pipe.vae._decode_without_taesd_preview(latents, *args, **kwargs)[0]
             # Compute reconstruction with the decoder
             # latents = latents.float()
-            res_taesd = compress(latents.mul(pipe.vae.config.scaling_factor)).mul_(2).sub_(1).cpu()
+            res_taesd = compress(latents.mul(pipe.vae.config.scaling_factor)).mul(2).sub(1)
             
             # Compute loss and backpropagate
-            loss = loss_fn(res_taesd, res_sd)
+            loss = loss_fn(res_taesd, res_sd.float())
             optimizer.zero_grad()
+            loss.cuda()
             loss.backward()
             optimizer.step()
             
@@ -105,6 +106,27 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
             torch.save(compress.state_dict(), f"{checkpoint_path}/compress_{optimizer.state_dict()['step']}.pth")
             return (res_taesd,)
         pipe.vae.decode = decode_latents_and_compare_speeds
+
+    print("debugging...")
+    # Simple training loop outside the pipeline
+    for _ in range(10):
+        print("creating latent")
+        latents = torch.randn(1, latent_channels, height, width, device='cuda', requires_grad=False)
+        print("decoding ")
+        res_sd = pipe.vae.decode(latents)[0]
+        print("compressing ")
+        res_taesd = compress(latents.mul(pipe.vae.config.scaling_factor)).mul(2).sub(1)
+        print("checking loss")
+        loss = loss_fn(res_taesd, res_sd.float())
+        optimizer.zero_grad()
+        print("backward")
+        loss.backward()
+        optimizer.step()
+
+        
+    import sys
+    sys.exit(1)
+
     
     add_final_speed_comparison(pipeline, instance, loss_fn, optimizer, checkpoint_path)
     # add_taesd_previewing(pipeline, instance)