Add comments to SD (#686)

* Add explanation for empty lambdas

* Fix my_unpickle if pytorch_lightning is installed

* oops

examples/stable_diffusion.py

@@ -163,7 +163,7 @@ class ResBlock:
     self.out_layers = [
       GroupNorm(32, out_channels),
       Tensor.silu,
-      lambda x: x,
+      lambda x: x,  # needed for weights loading code to work
       Conv2d(out_channels, out_channels, 3, padding=1)
     ]
     self.skip_connection = Conv2d(channels, out_channels, 1) if channels != out_channels else lambda x: x
@@ -213,7 +213,7 @@ class FeedForward:
   def __init__(self, dim, mult=4):
     self.net = [
       GEGLU(dim, dim*mult),
-      lambda x: x,
+      lambda x: x,  # needed for weights loading code to work
       Linear(dim*mult, dim)
     ]
 

extra/utils.py

@@ -74,8 +74,9 @@ def my_unpickle(fb0):
         elif name == "_rebuild_parameter":
           return HackParameter
       else:
+        if module.startswith('pytorch_lightning'): return Dummy
         try:
-          return pickle.Unpickler.find_class(self, module, name)
+          return super().find_class(module, name)
         except Exception:
           return Dummy