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>
2025-11-21 20:20:33 +01:00
parent 0999e5d29f
commit 9a964cff3c
4 changed files with 357 additions and 0 deletions
--- 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)
--- a/ai/FLUX_SETUP.md
+++ 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`
--- 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
--- a/ai/functions/flux_image_gen.py
+++ 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"
+            }]
+        })