File size: 2,865 Bytes
bd0bce1
 
944dd2b
26dad4e
bd0bce1
94bd22c
9d5a8fc
 
 
77d3539
9d5a8fc
2ced78d
bd0bce1
222ecc0
 
 
9d5a8fc
8c9a9e8
a51a106
 
 
9bf990d
 
 
bd0bce1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
513e628
944dd2b
26dad4e
 
bd0bce1
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
from .imagenhub_models import load_imagenhub_model
from .playground_api import load_playground_model
from .fal_api_models import load_fal_model
from .videogenhub_models import load_videogenhub_model


# IMAGE_GENERATION_MODELS = ['fal_LCM(v1.5/XL)_text2image','fal_SDXLTurbo_text2image','fal_SDXL_text2image', 'imagenhub_PixArtAlpha_generation', 'fal_PixArtSigma_text2image',
#                             'imagenhub_OpenJourney_generation','fal_SDXLLightning_text2image', 'fal_StableCascade_text2image',
#                             'playground_PlayGroundV2_generation', 'playground_PlayGroundV2.5_generation']
IMAGE_GENERATION_MODELS = ['imagenhub_SDXLTurbo_generation','imagenhub_SDXL_generation', 'imagenhub_PixArtAlpha_generation', 'imagenhub_PixArtSigma_generation',
                           'imagenhub_OpenJourney_generation','imagenhub_SDXLLightning_generation', 'imagenhub_StableCascade_generation',
                           'playground_PlayGroundV2.5_generation']  # 'playground_PlayGroundV2_generation'
IMAGE_EDITION_MODELS = ['imagenhub_CycleDiffusion_edition', 'imagenhub_Pix2PixZero_edition', 'imagenhub_Prompt2prompt_edition',
                        'imagenhub_SDEdit_edition', 'imagenhub_InstructPix2Pix_edition', 
                        'imagenhub_MagicBrush_edition', 'imagenhub_PNP_edition', 
                        'imagenhub_InfEdit_edition', 'imagenhub_CosXLEdit_edition']
VIDEO_GENERATION_MODELS = ['fal_AnimateDiff_text2video',
                           'fal_AnimateDiffTurbo_text2video',
                           'videogenhub_LaVie_generation', 
                           'videogenhub_VideoCrafter2_generation',
                           'videogenhub_ModelScope_generation',
                           'videogenhub_OpenSora_generation']
MUSEUM_UNSUPPORTED_MODELS = ['videogenhub_OpenSoraPlan_generation']
DESIRED_APPEAR_MODEL = ['videogenhub_T2VTurbo_generation','fal_StableVideoDiffusion_text2video']

def load_pipeline(model_name):
    """
    Load a model pipeline based on the model name
    Args:
        model_name (str): The name of the model to load, should be of the form {source}_{name}_{type}
        the source can be either imagenhub or playground
        the name is the name of the model used to load the model
        the type is the type of the model, either generation or edition
    """
    model_source, model_name, model_type = model_name.split("_")
    if model_source == "imagenhub":
        pipe = load_imagenhub_model(model_name, model_type)
    elif model_source == "playground":
        pipe = load_playground_model(model_name)
    elif model_source == "fal":
        pipe = load_fal_model(model_name, model_type)
    elif model_source == "videogenhub":
        pipe = load_videogenhub_model(model_name)
    else:
        raise ValueError(f"Model source {model_source} not supported")
    return pipe