Support segmind vega model.

This commit is contained in:
comfyanonymous
2023-12-12 19:09:53 -05:00
parent 824e4935f5
commit b454a67bb9
2 changed files with 18 additions and 2 deletions

View File

@@ -217,6 +217,16 @@ class SSD1B(SDXL):
"use_temporal_attention": False,
}
class Segmind_Vega(SDXL):
unet_config = {
"model_channels": 320,
"use_linear_in_transformer": True,
"transformer_depth": [0, 0, 1, 1, 2, 2],
"context_dim": 2048,
"adm_in_channels": 2816,
"use_temporal_attention": False,
}
class SVD_img2vid(supported_models_base.BASE):
unet_config = {
"model_channels": 320,
@@ -242,5 +252,5 @@ class SVD_img2vid(supported_models_base.BASE):
def clip_target(self):
return None
models = [SD15, SD20, SD21UnclipL, SD21UnclipH, SDXLRefiner, SDXL, SSD1B]
models = [SD15, SD20, SD21UnclipL, SD21UnclipH, SDXLRefiner, SDXL, SSD1B, Segmind_Vega]
models += [SVD_img2vid]