FlexTheAi/Flexstorydiff

1import yaml2import torch3from diffusers import StableDiffusionXLPipeline4from utils import PhotoMakerStableDiffusionXLPipeline5import os6 7def get_models_dict():8    # 打开并读取YAML文件9    with open('config/models.yaml', 'r') as stream:10        try:11            # 解析YAML文件内容12            data = yaml.safe_load(stream)13            14            # 此时 'data' 是一个Python字典，里面包含了YAML文件的所有数据15            print(data)16            return data17            18        except yaml.YAMLError as exc:19            # 如果在解析过程中发生了错误，打印异常信息20            print(exc)21 22def load_models(model_info,device,photomaker_path):23    path =  model_info["path"]24    single_files =  model_info["single_files"]25    use_safetensors = model_info["use_safetensors"]26    model_type = model_info["model_type"]27 28    if model_type == "original":29        if single_files:30            pipe = StableDiffusionXLPipeline.from_single_file(31                path, 32                torch_dtype=torch.float1633            )34        else:35            pipe = StableDiffusionXLPipeline.from_pretrained(path, torch_dtype=torch.float16, use_safetensors=use_safetensors)36        pipe = pipe.to(device)37    elif model_type == "Photomaker":38        if single_files:39            print("loading from a single_files")40            pipe = PhotoMakerStableDiffusionXLPipeline.from_single_file(41                path, 42                torch_dtype=torch.float1643            )44        else:45            pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(46                path, torch_dtype=torch.float16, use_safetensors=use_safetensors)47        pipe = pipe.to(device)48        pipe.load_photomaker_adapter(49            os.path.dirname(photomaker_path),50            subfolder="",51            weight_name=os.path.basename(photomaker_path),52            trigger_word="img"  # define the trigger word53        )54        pipe.fuse_lora()55    else:56        raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {model_type}")57    return pipe