Tool for calling stable diffusion from llm
| 12 | |
| 13 | |
| 14 | class StableDiffusionTool(GradioTool): |
| 15 | """Tool for calling stable diffusion from llm""" |
| 16 | |
| 17 | def __init__( |
| 18 | self, |
| 19 | name="StableDiffusion", |
| 20 | description=( |
| 21 | "An image generator. Use this to generate images based on " |
| 22 | "text input. Input should be a description of what the image should " |
| 23 | "look like. The output will be a path to an image file." |
| 24 | ), |
| 25 | src="gradio-client-demos/stable-diffusion", |
| 26 | hf_token=None, |
| 27 | duplicate=False, |
| 28 | ) -> None: |
| 29 | super().__init__(name, description, src, hf_token, duplicate) |
| 30 | |
| 31 | def create_job(self, query: str) -> Job: |
| 32 | return self.client.submit(query, "", 9, fn_index=1) |
| 33 | |
| 34 | def postprocess(self, output: Tuple[Any] | Any) -> str: |
| 35 | assert isinstance(output, str) |
| 36 | return [ |
| 37 | os.path.join(output, i) |
| 38 | for i in os.listdir(output) |
| 39 | if not i.endswith("json") |
| 40 | ][0] |
| 41 | |
| 42 | def _block_input(self, gr) -> "gr.components.Component": |
| 43 | return gr.Textbox() |
| 44 | |
| 45 | def _block_output(self, gr) -> "gr.components.Component": |
| 46 | return gr.Image() |