| 78 | print(output, end="\n\n") |
| 79 | |
| 80 | def build_tool(config) -> Tool: |
| 81 | tool = Tool( |
| 82 | "Image Generator", |
| 83 | "Tool that can generate image based on text description.", |
| 84 | name_for_model="Image Generator", |
| 85 | description_for_model=( |
| 86 | "Useful for when you need to generate an image." |
| 87 | "Input: A detailed text-2-image prompt describing an image" |
| 88 | "Output: the UUID of a generated image" |
| 89 | ), |
| 90 | logo_url="https://your-app-url.com/.well-known/logo.png", |
| 91 | contact_email="hello@contact.com", |
| 92 | legal_info_url="hello@legal.com" |
| 93 | ) |
| 94 | |
| 95 | model_name: ModelName = ModelName.DALL_E # choose model and image size? |
| 96 | size: Optional[str] = "512x512" |
| 97 | return_urls: Optional[bool] = False |
| 98 | steamship_api_key = os.environ.get('STEAMSHIP_API_KEY', '') |
| 99 | if steamship_api_key == '': |
| 100 | raise RuntimeError("STEAMSHIP_API_KEY is not provided. Please sign up for a free account at https://steamship.com/account/api, create a new API key, and add it to environment variables.") |
| 101 | |
| 102 | steamship = Steamship( |
| 103 | api_key=steamship_api_key, |
| 104 | ) |
| 105 | |
| 106 | @tool.get("/generate_image") |
| 107 | def generate_image(query : str): |
| 108 | '''Generate an image. |
| 109 | ''' |
| 110 | image_generator = steamship.use_plugin( |
| 111 | plugin_handle=model_name.value, config={"n": 1, "size": size} |
| 112 | ) |
| 113 | |
| 114 | task = image_generator.generate(text=query, append_output_to_file=True) |
| 115 | task.wait() |
| 116 | blocks = task.output.blocks |
| 117 | output_uiud = blocks[0].id |
| 118 | if len(blocks) > 0: |
| 119 | if return_urls: |
| 120 | output_uiud = make_image_public(steamship, blocks[0]) |
| 121 | # print image? |
| 122 | # show_output(output_uiud) |
| 123 | return output_uiud |
| 124 | raise RuntimeError("Tool unable to generate image!") |
| 125 | |
| 126 | return tool |