File tree Expand file tree Collapse file tree 2 files changed +8
-7
lines changed Expand file tree Collapse file tree 2 files changed +8
-7
lines changed Original file line number Diff line number Diff line change 41
41
STARTER_BUNDLES ,
42
42
STARTER_MODELS ,
43
43
StarterModel ,
44
+ StarterModelBundle ,
44
45
StarterModelWithoutDependencies ,
45
46
)
46
47
@@ -799,7 +800,7 @@ async def convert_model(
799
800
800
801
class StarterModelResponse (BaseModel ):
801
802
starter_models : list [StarterModel ]
802
- starter_bundles : dict [str , list [ StarterModel ] ]
803
+ starter_bundles : dict [str , StarterModelBundle ]
803
804
804
805
805
806
def get_is_installed (
@@ -833,7 +834,7 @@ async def get_starter_models() -> StarterModelResponse:
833
834
model .dependencies = missing_deps
834
835
835
836
for bundle in starter_bundles .values ():
836
- for model in bundle :
837
+ for model in bundle . models :
837
838
model .is_installed = get_is_installed (model , installed_models )
838
839
# Remove already-installed dependencies
839
840
missing_deps : list [StarterModelWithoutDependencies ] = []
Original file line number Diff line number Diff line change @@ -23,7 +23,7 @@ class StarterModel(StarterModelWithoutDependencies):
23
23
dependencies : Optional [list [StarterModelWithoutDependencies ]] = None
24
24
25
25
26
- class StarterModelBundles (BaseModel ):
26
+ class StarterModelBundle (BaseModel ):
27
27
name : str
28
28
models : list [StarterModel ]
29
29
@@ -778,10 +778,10 @@ class StarterModelBundles(BaseModel):
778
778
flux_fill ,
779
779
]
780
780
781
- STARTER_BUNDLES : dict [str , list [ StarterModel ] ] = {
782
- BaseModelType .StableDiffusion1 : sd1_bundle ,
783
- BaseModelType .StableDiffusionXL : sdxl_bundle ,
784
- BaseModelType .Flux : flux_bundle ,
781
+ STARTER_BUNDLES : dict [str , StarterModelBundle ] = {
782
+ BaseModelType .StableDiffusion1 : StarterModelBundle ( name = "Stable Diffusion 1.5" , models = sd1_bundle ) ,
783
+ BaseModelType .StableDiffusionXL : StarterModelBundle ( name = "SDXL" , models = sdxl_bundle ) ,
784
+ BaseModelType .Flux : StarterModelBundle ( name = "FLUX.1 dev" , models = flux_bundle ) ,
785
785
}
786
786
787
787
assert len (STARTER_MODELS ) == len ({m .source for m in STARTER_MODELS }), "Duplicate starter models"
You can’t perform that action at this time.
0 commit comments