修改 Stable Diffusion 使 api 接口增加模型参数

花果山总钻风2023-10-19 8:37

参考：https://zhuanlan.zhihu.com/p/644545784

1、修改 modules/api/models.py 中的 StableDiffusionTxt2ImgProcessingAPI 增加模型名称

python 复制代码

StableDiffusionTxt2ImgProcessingAPI = PydanticModelGenerator(
    "StableDiffusionProcessingTxt2Img",
    StableDiffusionProcessingTxt2Img,
    [
        {"key": "sampler_index", "type": str, "default": "Euler"},
        {"key": "script_name", "type": str, "default": None},
        {"key": "script_args", "type": list, "default": []},
        {"key": "send_images", "type": bool, "default": True},
        {"key": "save_images", "type": bool, "default": False},
        {"key": "alwayson_scripts", "type": dict, "default": {}},
        {"key": "model_name", "type": str, "default": None},
    ]
).generate_model()

2、修改 modules/processing.py 中的 StableDiffusionProcessingTxt2Img，增加模型名称接收

python 复制代码

@dataclass(repr=False)
class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
    enable_hr: bool = False
    denoising_strength: float = 0.75
    firstphase_width: int = 0
    firstphase_height: int = 0
    hr_scale: float = 2.0
    hr_upscaler: str = None
    hr_second_pass_steps: int = 0
    hr_resize_x: int = 0
    hr_resize_y: int = 0
    hr_checkpoint_name: str = None
    hr_sampler_name: str = None
    hr_prompt: str = ''
    hr_negative_prompt: str = ''
    model_name: str = None

3.修改 modules/api/api.py 中 text2imgapi 代码：

python 复制代码

......

from modules import sd_samplers, deepbooru, sd_hijack, images, scripts, ui, \
    postprocessing, errors, restart, shared_items, sd_models
from modules.api import models
from modules.shared import opts, models_path

......

def text2imgapi(self, txt2imgreq: models.StableDiffusionTxt2ImgProcessingAPI):
        ......
        model_name = txt2imgreq.model_name
        if model_name is None:
            raise HTTPException(status_code=404, detail="model_name not found")
        ......
        with self.queue_lock:
            checkpoint_info = sd_models.CheckpointInfo(os.path.join(models_path, 'Stable-diffusion', model_name))
            sd_models.reload_model_weights(info=checkpoint_info)
            with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p:
            ......