generate_flux_image
Generate images using the Flux model by submitting a text prompt via the Modal MCP Toolbox, designed for Python code execution and image creation in a sandbox environment.
Instructions
Let's you generate an image using the Flux model.
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| prompt | Yes | The prompt to generate an image for |
Implementation Reference
- src/modal_mcp_toolbox/flux.py:151-159 (handler)The handler function for the 'generate_flux_image' tool. It ensures the Modal app is deployed, calls the remote inference method on the Model class with the prompt, and converts the resulting image bytes to ImageContent with annotations.async def generate_flux_image(prompt: Annotated[str, Field(description="The prompt to generate an image for")], ctx: Context) -> ImageContent: """Let's you generate an image using the Flux model.""" await _ensure_app_deployment_is_up_to_date(ctx) cls = modal.Cls.from_name(app_name, Model.__name__) image_bytes = await cls().inference.remote.aio(prompt) image_content = Image(data=image_bytes, format=IMAGE_FORMAT).to_image_content() image_content.annotations = Annotations(audience=["user", "assistant"], priority=0.5) return image_content
- src/modal_mcp_toolbox/__init__.py:13-13 (registration)Registration of the 'generate_flux_image' tool with the FastMCP server.server.add_tool(generate_flux_image)
- src/modal_mcp_toolbox/flux.py:74-107 (helper)The Modal class 'Model' that loads the FluxPipeline, sets it up on GPU, and provides the 'inference' method to generate image bytes from a prompt using the Flux model.@app.cls( gpu="L40S", container_idle_timeout=5 * MINUTES, image=flux_image, volumes={ "/cache": modal.Volume.from_name("hf-hub-cache", create_if_missing=True), }, enable_memory_snapshot=True, ) class Model: @modal.enter(snap=True) def load(self): print("🔄 loading model...") pipe = FluxPipeline.from_pretrained(f"black-forest-labs/FLUX.1-{VARIANT}", torch_dtype=torch.bfloat16) self.pipe = _optimize(pipe) @modal.enter(snap=False) def setup(self): print("🔄 moving model to GPU...") self.pipe = self.pipe.to("cuda") @modal.method() def inference(self, prompt: str) -> bytes: print("🎨 generating image...") out = self.pipe( prompt, output_type="pil", num_inference_steps=NUM_INFERENCE_STEPS, ).images[0] byte_stream = BytesIO() out.save(byte_stream, format=IMAGE_FORMAT) return byte_stream.getvalue()
- Helper function called by the handler to ensure the Modal app deployment is up to date by checking version and deploying if necessary.async def _ensure_app_deployment_is_up_to_date(ctx: Context): try: fn = modal.Function.from_name(app_name, "get_version") remote_version = await fn.remote.aio() if remote_version != version("modal_mcp_toolbox"): await ctx.info("App is out of date. Deploying ...") logger.info("App is out of date. Deploying ...") deploy_app(app) except NotFoundError: await ctx.info("App not found. Deploying...") logger.info("App not found. Deploying...") deploy_app(app)
- Input schema definition via Annotated Field for the 'prompt' parameter.async def generate_flux_image(prompt: Annotated[str, Field(description="The prompt to generate an image for")], ctx: Context) -> ImageContent: