From 9a964cff3ceddf4d0d7b728a49993a688afbbe11 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20Kr=C3=BCger?= Date: Fri, 21 Nov 2025 20:20:33 +0100 Subject: [PATCH] feat: add Flux image generation function for Open WebUI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add flux_image_gen.py manifold function for Flux.1 Schnell - Auto-mount functions via Docker volume (./functions:/app/backend/data/functions:ro) - Add comprehensive setup guide in FLUX_SETUP.md - Update CLAUDE.md with Flux integration documentation - Infrastructure as code approach - no manual import needed 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- CLAUDE.md | 17 ++++ ai/FLUX_SETUP.md | 181 +++++++++++++++++++++++++++++++++ ai/compose.yaml | 1 + ai/functions/flux_image_gen.py | 158 ++++++++++++++++++++++++++++ 4 files changed, 357 insertions(+) create mode 100644 ai/FLUX_SETUP.md create mode 100644 ai/functions/flux_image_gen.py diff --git a/CLAUDE.md b/CLAUDE.md index 9eb37a3..38e3d04 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -476,11 +476,28 @@ AI infrastructure with Open WebUI, Crawl4AI, and dedicated PostgreSQL with pgvec 4. Use web search feature for current information 5. Integrate with n8n workflows for automation +**Flux Image Generation** (`functions/flux_image_gen.py`): +Open WebUI function for generating images via Flux.1 Schnell on RunPod GPU: +- Manifold function adds "Flux.1 Schnell (4-5s)" model to Open WebUI +- Routes requests through LiteLLM → Orchestrator → RunPod Flux +- Generates 1024x1024 images in 4-5 seconds +- Returns images as base64-encoded markdown +- Configuration via Valves (API base, timeout, default size) +- **Automatically loaded via Docker volume mount** (`./functions:/app/backend/data/functions:ro`) + +**Deployment**: +- Function file tracked in `ai/functions/` directory +- Automatically available after `pnpm arty up -d ai_webui` +- No manual import required - infrastructure as code + +See `ai/FLUX_SETUP.md` for detailed setup instructions and troubleshooting. + **Integration Points**: - **n8n**: Workflow automation with AI tasks (scraping, RAG ingestion, webhooks) - **Mattermost**: Can send AI-generated notifications via webhooks - **Crawl4AI**: Internal API for advanced web scraping - **Claude API**: Primary LLM provider via Anthropic +- **Flux via RunPod**: Image generation through orchestrator (GPU server) **Future Enhancements**: - GPU server integration (IONOS A10 planned) diff --git a/ai/FLUX_SETUP.md b/ai/FLUX_SETUP.md new file mode 100644 index 0000000..3eb7414 --- /dev/null +++ b/ai/FLUX_SETUP.md @@ -0,0 +1,181 @@ +# Flux Image Generation Setup for Open WebUI + +This guide explains how to add Flux.1 Schnell image generation to your Open WebUI installation. + +## Architecture + +``` +Open WebUI → flux_image_gen.py Function → LiteLLM (port 4000) → Orchestrator (RunPod port 9000) → Flux Model +``` + +## Installation + +### Automatic (via Docker Compose) + +The Flux function is **automatically loaded** via Docker volume mount. No manual upload needed! + +**How it works:** +- Function file: `ai/functions/flux_image_gen.py` +- Mounted to: `/app/backend/data/functions/` in the container (read-only) +- Open WebUI automatically discovers and loads functions from this directory on startup + +**To deploy:** +```bash +cd ~/Projects/docker-compose +pnpm arty up -d ai_webui # Restart Open WebUI to load function +``` + +### Verify Installation + +After restarting Open WebUI, the function should automatically appear in: +1. **Admin Settings → Functions**: Listed as "Flux Image Generator" +2. **Model dropdown**: "Flux.1 Schnell (4-5s)" available for selection + +If you don't see it: +```bash +# Check if function is mounted correctly +docker exec ai_webui ls -la /app/backend/data/functions/ + +# Check logs for any loading errors +docker logs ai_webui | grep -i flux +``` + +## Usage + +### Basic Image Generation + +1. **Select the Flux model:** + - In Open WebUI chat, select "Flux.1 Schnell (4-5s)" from the model dropdown + +2. **Send your prompt:** + ``` + A serene mountain landscape at sunset with vibrant colors + ``` + +3. **Wait for generation:** + - The function will call LiteLLM → Orchestrator → RunPod Flux + - Image appears in 4-5 seconds + +### Advanced Options + +The function supports custom sizes (configure in Valves): +- `1024x1024` (default, square) +- `1024x768` (landscape) +- `768x1024` (portrait) + +## Configuration + +### Valves (Customization) + +To customize function behavior: + +1. **Access Open WebUI**: + - Go to https://ai.pivoine.art + - Profile → Settings → Admin Settings → Functions + +2. **Find Flux Image Generator**: + - Click on "Flux Image Generator" in the functions list + - Go to "Valves" tab + +3. **Available Settings:** + - `LITELLM_API_BASE`: LiteLLM endpoint (default: `http://litellm:4000/v1`) + - `LITELLM_API_KEY`: API key (default: `dummy` - not needed for internal use) + - `DEFAULT_MODEL`: Model name (default: `flux-schnell`) + - `DEFAULT_SIZE`: Image dimensions (default: `1024x1024`) + - `TIMEOUT`: Request timeout in seconds (default: `120`) + +## Troubleshooting + +### Function not appearing in model list + +**Check:** +1. Function is enabled in Admin Settings → Functions +2. Function has no syntax errors (check logs) +3. Refresh browser cache (Ctrl+Shift+R) + +### Image generation fails + +**Check:** +1. LiteLLM is running: `docker ps | grep litellm` +2. LiteLLM can reach orchestrator: Check `docker logs ai_litellm` +3. Orchestrator is running on RunPod +4. Flux model is loaded: Check orchestrator logs + +**Test LiteLLM directly:** +```bash +curl -X POST http://localhost:4000/v1/images/generations \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "flux-schnell", + "prompt": "A test image", + "size": "1024x1024" + }' +``` + +### Timeout errors + +The default timeout is 120 seconds. If you're getting timeouts: + +1. **Increase timeout in Valves:** + - Set `TIMEOUT` to `180` or higher + +2. **Check Orchestrator status:** + - Flux model may still be loading (takes ~1 minute on first request) + +## Technical Details + +### How it Works + +1. **User sends prompt** in Open WebUI chat interface +2. **Function extracts prompt** from messages array +3. **Function calls LiteLLM** `/v1/images/generations` endpoint +4. **LiteLLM routes to Orchestrator** via config (`http://100.121.199.88:9000/v1`) +5. **Orchestrator loads Flux** on RunPod GPU (if not already running) +6. **Flux generates image** in 4-5 seconds +7. **Image returns as base64** through the chain +8. **Function displays image** as markdown in chat + +### Request Flow + +```json +// Function sends to LiteLLM: +{ + "model": "flux-schnell", + "prompt": "A serene mountain landscape", + "size": "1024x1024", + "n": 1, + "response_format": "b64_json" +} + +// LiteLLM response: +{ + "data": [{ + "b64_json": "iVBORw0KGgoAAAANSUhEUgAA..." + }] +} + +// Function converts to markdown: +![Generated Image](data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAA...) +``` + +## Limitations + +- **Single model**: Currently only Flux.1 Schnell is available +- **Sequential generation**: One image at a time (n=1) +- **Fixed format**: PNG format only +- **Orchestrator dependency**: Requires RunPod GPU server to be running + +## Future Enhancements + +Potential improvements: +- Multiple size presets in model dropdown +- Support for other Flux variants (Dev, Pro) +- Batch generation (n > 1) +- Image-to-image support +- Custom aspect ratios + +## Support + +- **Documentation**: `/home/valknar/Projects/docker-compose/CLAUDE.md` +- **RunPod README**: `/home/valknar/Projects/runpod/README.md` +- **LiteLLM Config**: `/home/valknar/Projects/docker-compose/ai/litellm-config.yaml` diff --git a/ai/compose.yaml b/ai/compose.yaml index 0a4ea06..a2318bd 100644 --- a/ai/compose.yaml +++ b/ai/compose.yaml @@ -66,6 +66,7 @@ services: volumes: - ai_webui_data:/app/backend/data + - ./functions:/app/backend/data/functions:ro depends_on: - ai_postgres - litellm diff --git a/ai/functions/flux_image_gen.py b/ai/functions/flux_image_gen.py new file mode 100644 index 0000000..fe2b392 --- /dev/null +++ b/ai/functions/flux_image_gen.py @@ -0,0 +1,158 @@ +""" +title: Flux Image Generator +author: Valknar +version: 1.0.0 +license: MIT +description: Generate images using Flux.1 Schnell via LiteLLM +requirements: requests, pydantic +""" + +import os +import base64 +import json +import requests +from typing import Generator +from pydantic import BaseModel, Field + + +class Pipe: + """ + Flux Image Generation Function for Open WebUI + Routes image generation requests to LiteLLM → Orchestrator → RunPod Flux + """ + + class Valves(BaseModel): + """Configuration valves for the image generation function""" + LITELLM_API_BASE: str = Field( + default="http://litellm:4000/v1", + description="LiteLLM API base URL" + ) + LITELLM_API_KEY: str = Field( + default="dummy", + description="LiteLLM API key (not required for internal use)" + ) + DEFAULT_MODEL: str = Field( + default="flux-schnell", + description="Default model to use for image generation" + ) + DEFAULT_SIZE: str = Field( + default="1024x1024", + description="Default image size" + ) + TIMEOUT: int = Field( + default=120, + description="Request timeout in seconds" + ) + + def __init__(self): + self.type = "manifold" + self.id = "flux_image_gen" + self.name = "Flux" + self.valves = self.Valves() + + def pipes(self): + """Return available models""" + return [ + { + "id": "flux-schnell", + "name": "Flux.1 Schnell (4-5s)" + } + ] + + def pipe(self, body: dict) -> Generator[str, None, None]: + """ + Generate images via LiteLLM endpoint + + Args: + body: Request body containing model, messages, etc. + + Yields: + JSON chunks with generated image data + """ + try: + # Extract the prompt from messages + messages = body.get("messages", []) + if not messages: + yield self._error_response("No messages provided") + return + + # Get the last user message as prompt + prompt = messages[-1].get("content", "") + if not prompt: + yield self._error_response("No prompt provided") + return + + # Prepare image generation request + image_request = { + "model": body.get("model", self.valves.DEFAULT_MODEL), + "prompt": prompt, + "size": body.get("size", self.valves.DEFAULT_SIZE), + "n": 1, + "response_format": "b64_json" + } + + # Call LiteLLM images endpoint + response = requests.post( + f"{self.valves.LITELLM_API_BASE}/images/generations", + json=image_request, + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {self.valves.LITELLM_API_KEY}" + }, + timeout=self.valves.TIMEOUT + ) + + if response.status_code != 200: + yield self._error_response( + f"Image generation failed: {response.status_code} - {response.text}" + ) + return + + # Parse response + result = response.json() + + # Check if we got image data + if "data" not in result or len(result["data"]) == 0: + yield self._error_response("No image data in response") + return + + # Get base64 image data + image_data = result["data"][0].get("b64_json") + if not image_data: + yield self._error_response("No base64 image data in response") + return + + # Return image as markdown + image_markdown = f"![Generated Image](data:image/png;base64,{image_data})\n\n**Prompt:** {prompt}" + + # Yield final response + yield json.dumps({ + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": image_markdown + }, + "finish_reason": "stop" + }] + }) + + except requests.Timeout: + yield self._error_response(f"Request timed out after {self.valves.TIMEOUT}s") + except requests.RequestException as e: + yield self._error_response(f"Request failed: {str(e)}") + except Exception as e: + yield self._error_response(f"Unexpected error: {str(e)}") + + def _error_response(self, error_message: str) -> str: + """Generate error response in OpenAI format""" + return json.dumps({ + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": f"Error: {error_message}" + }, + "finish_reason": "stop" + }] + })