| from modules import scripts_postprocessing, ui_components, deepbooru, shared | |
| import gradio as gr | |
| class ScriptPostprocessingCeption(scripts_postprocessing.ScriptPostprocessing): | |
| name = "Caption" | |
| order = 4040 | |
| def ui(self): | |
| with ui_components.InputAccordion(False, label="Caption") as enable: | |
| option = gr.CheckboxGroup(value=["Deepbooru"], choices=["Deepbooru", "BLIP"], show_label=False) | |
| return { | |
| "enable": enable, | |
| "option": option, | |
| } | |
| def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, option): | |
| if not enable: | |
| return | |
| captions = [pp.caption] | |
| if "Deepbooru" in option: | |
| captions.append(deepbooru.model.tag(pp.image)) | |
| if "BLIP" in option: | |
| captions.append(shared.interrogator.interrogate(pp.image.convert("RGB"))) | |
| pp.caption = ", ".join([x for x in captions if x]) | |