Merge 2336124dd4 into 53764fe648

README.md

@@ -47,7 +47,7 @@ For more information, be sure to check out our [Open WebUI Documentation](https:
 
 - 🌐 **Web Browsing Capability**: Seamlessly integrate websites into your chat experience using the `#` command followed by a URL. This feature allows you to incorporate web content directly into your conversations, enhancing the richness and depth of your interactions.
 
-- 🎨 **Image Generation Integration**: Seamlessly incorporate image generation capabilities using options such as AUTOMATIC1111 API or ComfyUI (local), and OpenAI's DALL-E (external), enriching your chat experience with dynamic visual content.
+- 🎨 **Image Generation Integration**: Seamlessly incorporate image generation capabilities using options such as AUTOMATIC1111 API, ComfyUI (local), Replicate, and OpenAI's DALL-E (external), enriching your chat experience with dynamic visual content.
 
 - ⚙️ **Many Models Conversations**: Effortlessly engage with various models simultaneously, harnessing their unique strengths for optimal responses. Enhance your experience by leveraging a diverse set of models in parallel.
 

backend/open_webui/config.py

@@ -2982,3 +2982,13 @@ LDAP_VALIDATE_CERT = PersistentConfig(
 LDAP_CIPHERS = PersistentConfig(
     "LDAP_CIPHERS", "ldap.server.ciphers", os.environ.get("LDAP_CIPHERS", "ALL")
 )
+
+####################################
+# REPLICATE
+####################################
+
+REPLICATE_API_TOKEN = PersistentConfig(
+    "REPLICATE_API_TOKEN",
+    "image_generation.replicate.api_token",
+    os.getenv("REPLICATE_API_TOKEN", ""),
+)

backend/open_webui/main.py

@@ -341,6 +341,8 @@ from open_webui.config import (
     LDAP_CA_CERT_FILE,
     LDAP_VALIDATE_CERT,
     LDAP_CIPHERS,
+    # Replicate
+    REPLICATE_API_TOKEN,
     # Misc
     ENV,
     CACHE_DIR,
@@ -642,6 +644,7 @@ app.state.config.LDAP_CA_CERT_FILE = LDAP_CA_CERT_FILE
 app.state.config.LDAP_VALIDATE_CERT = LDAP_VALIDATE_CERT
 app.state.config.LDAP_CIPHERS = LDAP_CIPHERS
 
+app.state.config.REPLICATE_API_TOKEN = REPLICATE_API_TOKEN
 
 app.state.AUTH_TRUSTED_EMAIL_HEADER = WEBUI_AUTH_TRUSTED_EMAIL_HEADER
 app.state.AUTH_TRUSTED_NAME_HEADER = WEBUI_AUTH_TRUSTED_NAME_HEADER

backend/test_replicate_unit.py

@@ -0,0 +1,89 @@
+"""
+Simple unit tests for Replicate image generation functionality.
+Run with: python -m pytest test_replicate_unit.py -v
+"""
+import pytest
+from unittest.mock import patch, MagicMock
+import sys
+import os
+
+# Add the open_webui module to path for imports
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'open_webui'))
+
+from open_webui.routers.images import get_replicate_models
+
+
+class TestReplicateModels:
+    """Test Replicate model fetching functionality"""
+    
+    def test_get_replicate_models_no_token(self):
+        """Test that cached models are returned when no API token is provided"""
+        result = get_replicate_models("")
+        
+        assert len(result) == 14
+        assert all("id" in model for model in result)
+        assert all("name" in model for model in result)
+        assert all("description" in model for model in result)
+        
+        # Check that default model is included
+        model_ids = [model["id"] for model in result]
+        assert "black-forest-labs/flux-1.1-pro-ultra" in model_ids
+        assert "black-forest-labs/flux-1.1-pro" in model_ids
+
+    def test_get_replicate_models_with_token_fallback(self):
+        """Test that cached models are returned when API fails"""
+        with patch('open_webui.routers.images.replicate') as mock_replicate:
+            # Mock API failure
+            mock_replicate.models.get.side_effect = Exception("API Error")
+            
+            result = get_replicate_models("test_token")
+            
+            # Should still return cached models
+            assert len(result) == 14
+            assert result[0]["id"] == "black-forest-labs/flux-1.1-pro-ultra"
+
+    def test_get_replicate_models_cached_structure(self):
+        """Test that cached models have the correct structure"""
+        result = get_replicate_models("")
+        
+        for model in result:
+            assert "id" in model
+            assert "name" in model  
+            assert "description" in model
+            assert isinstance(model["id"], str)
+            assert isinstance(model["name"], str)
+            assert isinstance(model["description"], str)
+            assert len(model["id"]) > 0
+            assert len(model["name"]) > 0
+
+    def test_replicate_models_include_popular_options(self):
+        """Test that cached models include popular/expected models"""
+        result = get_replicate_models("")
+        model_ids = [model["id"] for model in result]
+        
+        # Check for key models that should be included
+        expected_models = [
+            "black-forest-labs/flux-1.1-pro-ultra",
+            "black-forest-labs/flux-1.1-pro", 
+            "black-forest-labs/flux-schnell",
+            "stability-ai/stable-diffusion-3.5-large",
+            "stability-ai/sdxl"
+        ]
+        
+        for expected in expected_models:
+            assert expected in model_ids, f"Expected model {expected} not found in cached models"
+
+    def test_get_replicate_models_return_type(self):
+        """Test that the function returns a list of dictionaries"""
+        result = get_replicate_models("")
+        
+        assert isinstance(result, list)
+        assert len(result) > 0
+        
+        for model in result:
+            assert isinstance(model, dict)
+
+
+if __name__ == "__main__":
+    # Allow running tests directly
+    pytest.main([__file__, "-v"]) 

pyproject.toml

@@ -134,6 +134,7 @@ dependencies = [
 
     "moto[s3]>=5.0.26",
 
+    "replicate",
 ]
 readme = "README.md"
 requires-python = ">= 3.11, < 3.13.0a1"

src/lib/components/admin/Settings/Images.svelte

@@ -117,6 +117,11 @@
 
 		if (res) {
 			config = res;
+			if (!config.replicate) {
+				config.replicate = {
+					REPLICATE_API_TOKEN: ''
+				};
+			}
 		}
 
 		if (config.enabled) {
@@ -184,6 +189,11 @@
 
 			if (res) {
 				config = res;
+				if (!config.replicate) {
+					config.replicate = {
+						REPLICATE_API_TOKEN: ''
+					};
+				}
 			}
 
 			if (config.enabled) {
@@ -302,6 +312,7 @@
 							<option value="comfyui">{$i18n.t('ComfyUI')}</option>
 							<option value="automatic1111">{$i18n.t('Automatic1111')}</option>
 							<option value="gemini">{$i18n.t('Gemini')}</option>
+							<option value="replicate">{$i18n.t('Replicate')}</option>
 						</select>
 					</div>
 				</div>
@@ -631,6 +642,17 @@
 							/>
 						</div>
 					</div>
+				{:else if config?.engine === 'replicate'}
+					<div>
+						<div class=" mb-1.5 text-sm font-medium">{$i18n.t('Replicate API Config')}</div>
+
+						<div class="flex gap-2 mb-1">
+							<SensitiveInput
+								placeholder={$i18n.t('API Key')}
+								bind:value={config.replicate.REPLICATE_API_TOKEN}
+							/>
+						</div>
+					</div>
 				{/if}
 			</div>