Working multi gpu training. Still need a lot of tweaks and testing.

toolkit/accelerator.py

@@ -0,0 +1,17 @@
+from accelerate import Accelerator
+from diffusers.utils.torch_utils import is_compiled_module
+
+global_accelerator = None
+
+
+def get_accelerator() -> Accelerator:
+    global global_accelerator
+    if global_accelerator is None:
+        global_accelerator = Accelerator()
+    return global_accelerator
+
+def unwrap_model(model):
+    accelerator = get_accelerator()
+    model = accelerator.unwrap_model(model)
+    model = model._orig_mod if is_compiled_module(model) else model
+    return model