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:
17
CLAUDE.md
17
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
|
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
181
ai/FLUX_SETUP.md
Normal 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:
|
||||||
|

|
||||||
|
```
|
||||||
|
|
||||||
|
## 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`
|
||||||
@@ -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
|
||||||
|
|||||||
158
ai/functions/flux_image_gen.py
Normal file
158
ai/functions/flux_image_gen.py
Normal 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"\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"
|
||||||
|
}]
|
||||||
|
})
|
||||||
Reference in New Issue
Block a user