feat: add Flux image generation function for Open WebUI

- 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 <noreply@anthropic.com>
This commit is contained in:
2025-11-21 20:20:33 +01:00
parent 0999e5d29f
commit 9a964cff3c
4 changed files with 357 additions and 0 deletions

View File

@@ -476,11 +476,28 @@ AI infrastructure with Open WebUI, Crawl4AI, and dedicated PostgreSQL with pgvec
4. Use web search feature for current information 4. Use web search feature for current information
5. Integrate with n8n workflows for automation 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**: **Integration Points**:
- **n8n**: Workflow automation with AI tasks (scraping, RAG ingestion, webhooks) - **n8n**: Workflow automation with AI tasks (scraping, RAG ingestion, webhooks)
- **Mattermost**: Can send AI-generated notifications via webhooks - **Mattermost**: Can send AI-generated notifications via webhooks
- **Crawl4AI**: Internal API for advanced web scraping - **Crawl4AI**: Internal API for advanced web scraping
- **Claude API**: Primary LLM provider via Anthropic - **Claude API**: Primary LLM provider via Anthropic
- **Flux via RunPod**: Image generation through orchestrator (GPU server)
**Future Enhancements**: **Future Enhancements**:
- GPU server integration (IONOS A10 planned) - GPU server integration (IONOS A10 planned)

181
ai/FLUX_SETUP.md Normal file
View File

@@ -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`

View File

@@ -66,6 +66,7 @@ services:
volumes: volumes:
- ai_webui_data:/app/backend/data - ai_webui_data:/app/backend/data
- ./functions:/app/backend/data/functions:ro
depends_on: depends_on:
- ai_postgres - ai_postgres
- litellm - litellm

View File

@@ -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"
}]
})