Spaces:
Sleeping
Sleeping
| # tools/text_to_image.py | |
| import os | |
| from huggingface_hub import InferenceClient | |
| from smolagents import Tool | |
| class TextToImageTool(Tool): | |
| name = "image_generator" | |
| description = "Generate an image from a text prompt." | |
| inputs = { | |
| "prompt": { | |
| "type": "string", | |
| "description": "Text prompt for the image. Add style hints like photorealistic, high-res, etc.", | |
| } | |
| } | |
| output_type = "image" | |
| def __init__(self, model_id: str = "stabilityai/stable-diffusion-xl-base-1.0"): | |
| super().__init__() | |
| token = ( | |
| os.getenv("HF_TOKEN") | |
| or os.getenv("HUGGINGFACEHUB_API_TOKEN") | |
| or os.getenv("HF_API_TOKEN") | |
| ) | |
| self.client = InferenceClient(model=model_id, token=token) | |
| def forward(self, prompt: str): | |
| return self.client.text_to_image(prompt) |