Added cogview4. Loss still needs work.

2025-03-04 18:43:52 -07:00
parent c57434ad7b
commit 6f6fb90812
12 changed files with 661 additions and 152 deletions
--- a/toolkit/util/get_model.py
+++ b/toolkit/util/get_model.py
@@ -5,5 +5,8 @@ def get_model_class(config: ModelConfig):
    if config.arch == "wan21":
        from toolkit.models.wan21 import Wan21
        return Wan21
+    elif config.arch == "cogview4":
+        from toolkit.models.cogview4 import CogView4
+        return CogView4
    else:
        return StableDiffusion