feat: replace local image provider with ComfyUI configuration in environment variables and user settings

This commit is contained in:
shiva raj badu 2025-12-18 22:42:53 +05:45
parent f36fc7f4bb
commit ea8db1bfdd
No known key found for this signature in database
16 changed files with 94 additions and 107 deletions

View file

@ -30,8 +30,8 @@ services:
- WEB_GROUNDING=${WEB_GROUNDING}
- DATABASE_URL=${DATABASE_URL}
- DISABLE_ANONYMOUS_TRACKING=${DISABLE_ANONYMOUS_TRACKING}
- LOCAL_IMAGE_URL=${LOCAL_IMAGE_URL}
- LOCAL_IMAGE_WORKFLOW=${LOCAL_IMAGE_WORKFLOW}
- COMFYUI_URL=${COMFYUI_URL}
- COMFYUI_WORKFLOW=${COMFYUI_WORKFLOW}
production-gpu:
# image: ghcr.io/presenton/presenton:latest
@ -71,8 +71,8 @@ services:
- WEB_GROUNDING=${WEB_GROUNDING}
- DATABASE_URL=${DATABASE_URL}
- DISABLE_ANONYMOUS_TRACKING=${DISABLE_ANONYMOUS_TRACKING}
- LOCAL_IMAGE_URL=${LOCAL_IMAGE_URL}
- LOCAL_IMAGE_WORKFLOW=${LOCAL_IMAGE_WORKFLOW}
- COMFYUI_URL=${COMFYUI_URL}
- COMFYUI_WORKFLOW=${COMFYUI_WORKFLOW}
development:
build:
@ -104,8 +104,8 @@ services:
- WEB_GROUNDING=${WEB_GROUNDING}
- DATABASE_URL=${DATABASE_URL}
- DISABLE_ANONYMOUS_TRACKING=${DISABLE_ANONYMOUS_TRACKING}
- LOCAL_IMAGE_URL=${LOCAL_IMAGE_URL}
- LOCAL_IMAGE_WORKFLOW=${LOCAL_IMAGE_WORKFLOW}
- COMFYUI_URL=${COMFYUI_URL}
- COMFYUI_WORKFLOW=${COMFYUI_WORKFLOW}
development-gpu:
build:
@ -144,5 +144,5 @@ services:
- WEB_GROUNDING=${WEB_GROUNDING}
- DATABASE_URL=${DATABASE_URL}
- DISABLE_ANONYMOUS_TRACKING=${DISABLE_ANONYMOUS_TRACKING}
- LOCAL_IMAGE_URL=${LOCAL_IMAGE_URL}
- LOCAL_IMAGE_WORKFLOW=${LOCAL_IMAGE_WORKFLOW}
- COMFYUI_URL=${COMFYUI_URL}
- COMFYUI_WORKFLOW=${COMFYUI_WORKFLOW}

View file

@ -5,4 +5,4 @@ class ImageProvider(Enum):
PIXABAY = "pixabay"
GEMINI_FLASH = "gemini_flash"
DALLE3 = "dall-e-3"
LOCAL = "local" # Local image generation (Stable Diffusion, FLUX, ComfyUI, Fooocus, etc.)
COMFYUI = "comfyui"

View file

@ -32,9 +32,9 @@ class UserConfig(BaseModel):
PEXELS_API_KEY: Optional[str] = None
PIXABAY_API_KEY: Optional[str] = None
# Local Image Generation (ComfyUI)
LOCAL_IMAGE_URL: Optional[str] = None
LOCAL_IMAGE_WORKFLOW: Optional[str] = None # ComfyUI workflow JSON
# ComfyUI
COMFYUI_URL: Optional[str] = None
COMFYUI_WORKFLOW: Optional[str] = None
# Reasoning
TOOL_CALLS: Optional[bool] = None

View file

@ -11,15 +11,15 @@ from models.sql.image_asset import ImageAsset
from utils.download_helpers import download_file
from utils.get_env import get_pexels_api_key_env
from utils.get_env import get_pixabay_api_key_env
from utils.get_env import get_local_image_url_env
from utils.get_env import get_local_image_workflow_env
from utils.get_env import get_comfyui_url_env
from utils.get_env import get_comfyui_workflow_env
from utils.image_provider import (
is_image_generation_disabled,
is_pixels_selected,
is_pixabay_selected,
is_gemini_flash_selected,
is_dalle3_selected,
is_local_selected,
is_comfyui_selected,
)
import uuid
@ -42,8 +42,8 @@ class ImageGenerationService:
return self.generate_image_google
elif is_dalle3_selected():
return self.generate_image_openai
elif is_local_selected():
return self.generate_image_local
elif is_comfyui_selected():
return self.generate_image_comfyui
return None
def is_stock_provider_selected(self):
@ -145,13 +145,13 @@ class ImageGenerationService:
image_url = data["hits"][0]["largeImageURL"]
return image_url
async def generate_image_local(self, prompt: str, output_directory: str) -> str:
async def generate_image_comfyui(self, prompt: str, output_directory: str) -> str:
"""
Generate image using ComfyUI workflow API.
User provides:
- LOCAL_IMAGE_URL: ComfyUI server URL (e.g., http://192.168.1.7:8188)
- LOCAL_IMAGE_WORKFLOW: Workflow JSON exported from ComfyUI
- COMFYUI_URL: ComfyUI server URL (e.g., http://192.168.1.7:8188)
- COMFYUI_WORKFLOW: Workflow JSON exported from ComfyUI
The workflow should have a CLIPTextEncode node with "Positive" in the title
where the prompt will be injected.
@ -163,14 +163,14 @@ class ImageGenerationService:
Returns:
Path to the generated image file
"""
comfyui_url = get_local_image_url_env()
workflow_json = get_local_image_workflow_env()
comfyui_url = get_comfyui_url_env()
workflow_json = get_comfyui_workflow_env()
if not comfyui_url:
raise ValueError("LOCAL_IMAGE_URL environment variable is not set")
raise ValueError("COMFYUI_URL environment variable is not set")
if not workflow_json:
raise ValueError("LOCAL_IMAGE_WORKFLOW environment variable is not set. Please provide a ComfyUI workflow JSON.")
raise ValueError("COMFYUI_WORKFLOW environment variable is not set. Please provide a ComfyUI workflow JSON.")
# Ensure URL doesn't have trailing slash
comfyui_url = comfyui_url.rstrip("/")
@ -200,44 +200,22 @@ class ImageGenerationService:
def _inject_prompt_into_workflow(self, workflow: dict, prompt: str) -> dict:
"""
Find the positive prompt node in the workflow and inject the prompt text.
Looks for CLIPTextEncode nodes with 'Positive' in the title.
Find the prompt node in the workflow and inject the prompt text.
Looks for a node with title 'Input Prompt' (case-insensitive).
User must rename their prompt node to 'Input Prompt' in ComfyUI.
"""
prompt_injected = False
for node_id, node_data in workflow.items():
# Check if this is a CLIPTextEncode node
if node_data.get("class_type") == "CLIPTextEncode":
meta = node_data.get("_meta", {})
title = meta.get("title", "").lower()
# Check if it's a positive prompt node
if "positive" in title:
if "inputs" in node_data and "text" in node_data["inputs"]:
node_data["inputs"]["text"] = prompt
prompt_injected = True
print(f"Injected prompt into node {node_id}: {title}")
break
meta = node_data.get("_meta", {})
title = meta.get("title", "").lower()
if title == "input prompt":
if "inputs" in node_data and "text" in node_data["inputs"]:
node_data["inputs"]["text"] = prompt
print(f"Injected prompt into node {node_id}: {meta.get('title', '')}")
return workflow
if not prompt_injected:
# Fallback: try to find any CLIPTextEncode node with text input
for node_id, node_data in workflow.items():
if node_data.get("class_type") == "CLIPTextEncode":
if "inputs" in node_data and "text" in node_data["inputs"]:
# Skip if it looks like a negative prompt
meta = node_data.get("_meta", {})
title = meta.get("title", "").lower()
if "negative" in title:
continue
node_data["inputs"]["text"] = prompt
prompt_injected = True
print(f"Injected prompt into node {node_id} (fallback)")
break
if not prompt_injected:
raise ValueError("Could not find a positive prompt node (CLIPTextEncode) in the workflow")
return workflow
raise ValueError("Could not find a node with title 'Input Prompt' in the workflow. Please rename your prompt node to 'Input Prompt' in ComfyUI.")
async def _submit_comfyui_workflow(
self, session: aiohttp.ClientSession, comfyui_url: str, workflow: dict

View file

@ -101,9 +101,9 @@ def get_web_grounding_env():
return os.getenv("WEB_GROUNDING")
def get_local_image_url_env():
return os.getenv("LOCAL_IMAGE_URL")
def get_comfyui_url_env():
return os.getenv("COMFYUI_URL")
def get_local_image_workflow_env():
return os.getenv("LOCAL_IMAGE_WORKFLOW")
def get_comfyui_workflow_env():
return os.getenv("COMFYUI_WORKFLOW")

View file

@ -1,9 +1,9 @@
from enums.image_provider import ImageProvider
from utils.get_env import (
get_comfyui_url_env,
get_disable_image_generation_env,
get_google_api_key_env,
get_image_provider_env,
get_local_image_url_env,
get_openai_api_key_env,
get_pexels_api_key_env,
get_pixabay_api_key_env,
@ -31,8 +31,8 @@ def is_dalle3_selected() -> bool:
return ImageProvider.DALLE3 == get_selected_image_provider()
def is_local_selected() -> bool:
return ImageProvider.LOCAL == get_selected_image_provider()
def is_comfyui_selected() -> bool:
return ImageProvider.COMFYUI == get_selected_image_provider()
def get_selected_image_provider() -> ImageProvider | None:
@ -57,7 +57,7 @@ def get_image_provider_api_key() -> str:
return get_google_api_key_env()
elif selected_image_provider == ImageProvider.DALLE3:
return get_openai_api_key_env()
elif selected_image_provider == ImageProvider.LOCAL:
return get_local_image_url_env() # Returns URL instead of API key
elif selected_image_provider == ImageProvider.COMFYUI:
return get_comfyui_url_env() # Returns URL instead of API key
else:
raise ValueError(f"Invalid image provider: {selected_image_provider}")

View file

@ -16,6 +16,8 @@ from utils.get_env import (
get_openai_model_env,
get_pixabay_api_key_env,
get_pexels_api_key_env,
get_comfyui_url_env,
get_comfyui_workflow_env,
)
from utils.get_env import get_google_api_key_env
from utils.get_env import get_ollama_model_env
@ -135,3 +137,10 @@ async def check_llm_and_image_provider_api_or_model_availability():
openai_api_key = get_openai_api_key_env()
if not openai_api_key:
raise Exception("OPENAI_API_KEY must be provided")
elif selected_image_provider == ImageProvider.COMFYUI:
comfyui_url = get_comfyui_url_env()
if not comfyui_url:
raise Exception("COMFYUI_URL must be provided")
workflow_json = get_comfyui_workflow_env()
if not workflow_json:
raise Exception("COMFYUI_WORKFLOW must be provided")

View file

@ -89,9 +89,9 @@ def set_web_grounding_env(value):
os.environ["WEB_GROUNDING"] = value
def set_local_image_url_env(value):
os.environ["LOCAL_IMAGE_URL"] = value
def set_comfyui_url_env(value):
os.environ["COMFYUI_URL"] = value
def set_local_image_workflow_env(value):
os.environ["LOCAL_IMAGE_WORKFLOW"] = value
def set_comfyui_workflow_env(value):
os.environ["COMFYUI_WORKFLOW"] = value

View file

@ -5,6 +5,8 @@ from models.user_config import UserConfig
from utils.get_env import (
get_anthropic_api_key_env,
get_anthropic_model_env,
get_comfyui_url_env,
get_comfyui_workflow_env,
get_custom_llm_api_key_env,
get_custom_llm_url_env,
get_custom_model_env,
@ -13,8 +15,6 @@ from utils.get_env import (
get_google_api_key_env,
get_google_model_env,
get_llm_provider_env,
get_local_image_url_env,
get_local_image_workflow_env,
get_ollama_model_env,
get_ollama_url_env,
get_openai_api_key_env,
@ -31,6 +31,8 @@ from utils.parsers import parse_bool_or_none
from utils.set_env import (
set_anthropic_api_key_env,
set_anthropic_model_env,
set_comfyui_url_env,
set_comfyui_workflow_env,
set_custom_llm_api_key_env,
set_custom_llm_url_env,
set_custom_model_env,
@ -40,8 +42,6 @@ from utils.set_env import (
set_google_api_key_env,
set_google_model_env,
set_llm_provider_env,
set_local_image_url_env,
set_local_image_workflow_env,
set_ollama_model_env,
set_ollama_url_env,
set_openai_api_key_env,
@ -89,8 +89,8 @@ def get_user_config():
),
PIXABAY_API_KEY=existing_config.PIXABAY_API_KEY or get_pixabay_api_key_env(),
PEXELS_API_KEY=existing_config.PEXELS_API_KEY or get_pexels_api_key_env(),
LOCAL_IMAGE_URL=existing_config.LOCAL_IMAGE_URL or get_local_image_url_env(),
LOCAL_IMAGE_WORKFLOW=existing_config.LOCAL_IMAGE_WORKFLOW or get_local_image_workflow_env(),
COMFYUI_URL=existing_config.COMFYUI_URL or get_comfyui_url_env(),
COMFYUI_WORKFLOW=existing_config.COMFYUI_WORKFLOW or get_comfyui_workflow_env(),
TOOL_CALLS=(
existing_config.TOOL_CALLS
if existing_config.TOOL_CALLS is not None
@ -148,10 +148,10 @@ def update_env_with_user_config():
set_pixabay_api_key_env(user_config.PIXABAY_API_KEY)
if user_config.PEXELS_API_KEY:
set_pexels_api_key_env(user_config.PEXELS_API_KEY)
if user_config.LOCAL_IMAGE_URL:
set_local_image_url_env(user_config.LOCAL_IMAGE_URL)
if user_config.LOCAL_IMAGE_WORKFLOW:
set_local_image_workflow_env(user_config.LOCAL_IMAGE_WORKFLOW)
if user_config.COMFYUI_URL:
set_comfyui_url_env(user_config.COMFYUI_URL)
if user_config.COMFYUI_WORKFLOW:
set_comfyui_workflow_env(user_config.COMFYUI_WORKFLOW)
if user_config.TOOL_CALLS is not None:
set_tool_calls_env(str(user_config.TOOL_CALLS))
if user_config.DISABLE_THINKING is not None:

View file

@ -64,8 +64,8 @@ export async function POST(request: Request) {
userConfig.PIXABAY_API_KEY || existingConfig.PIXABAY_API_KEY,
IMAGE_PROVIDER: userConfig.IMAGE_PROVIDER || existingConfig.IMAGE_PROVIDER,
PEXELS_API_KEY: userConfig.PEXELS_API_KEY || existingConfig.PEXELS_API_KEY,
LOCAL_IMAGE_URL: userConfig.LOCAL_IMAGE_URL || existingConfig.LOCAL_IMAGE_URL,
LOCAL_IMAGE_WORKFLOW: userConfig.LOCAL_IMAGE_WORKFLOW || existingConfig.LOCAL_IMAGE_WORKFLOW,
COMFYUI_URL: userConfig.COMFYUI_URL || existingConfig.COMFYUI_URL,
COMFYUI_WORKFLOW: userConfig.COMFYUI_WORKFLOW || existingConfig.COMFYUI_WORKFLOW,
TOOL_CALLS:
userConfig.TOOL_CALLS === undefined
? existingConfig.TOOL_CALLS

View file

@ -83,8 +83,8 @@ export default function LLMProviderSelection({
const needsOllamaUrl = (llmConfig.LLM === "ollama" && !llmConfig.OLLAMA_URL);
const needsComfyUIConfig = !llmConfig.DISABLE_IMAGE_GENERATION &&
llmConfig.IMAGE_PROVIDER === "local" &&
(!llmConfig.LOCAL_IMAGE_URL || !llmConfig.LOCAL_IMAGE_WORKFLOW);
llmConfig.IMAGE_PROVIDER === "comfyui" &&
(!llmConfig.COMFYUI_URL || !llmConfig.COMFYUI_WORKFLOW);
setButtonState({
isLoading: false,
@ -341,7 +341,7 @@ export default function LLMProviderSelection({
}
// Show ComfyUI configuration
if (provider.value === "local") {
if (provider.value === "comfyui") {
return (
<div className="mb-8 space-y-4">
<div>
@ -353,9 +353,9 @@ export default function LLMProviderSelection({
type="text"
placeholder="http://192.168.1.7:8188"
className="w-full px-4 py-2.5 outline-none border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500/20 focus:border-blue-500 transition-colors"
value={llmConfig.LOCAL_IMAGE_URL || ""}
value={llmConfig.COMFYUI_URL || ""}
onChange={(e) => {
input_field_changed(e.target.value, "local_image_url");
input_field_changed(e.target.value, "comfyui_url");
}}
/>
</div>
@ -370,18 +370,18 @@ export default function LLMProviderSelection({
</label>
<div className="relative">
<textarea
placeholder='Paste your ComfyUI workflow JSON here (export via "Save (API Format)" in ComfyUI)'
placeholder='Paste your ComfyUI workflow JSON here (export via "Export (API)" in ComfyUI)'
className="w-full px-4 py-2.5 outline-none border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500/20 focus:border-blue-500 transition-colors font-mono text-xs"
rows={6}
value={llmConfig.LOCAL_IMAGE_WORKFLOW || ""}
value={llmConfig.COMFYUI_WORKFLOW || ""}
onChange={(e) => {
input_field_changed(e.target.value, "local_image_workflow");
input_field_changed(e.target.value, "comfyui_workflow");
}}
/>
</div>
<p className="mt-2 text-sm text-gray-500">
Export your workflow from ComfyUI using &quot;Save (API Format)&quot; and paste the JSON here.
The positive prompt node (CLIPTextEncode) will be automatically updated.
Export your workflow from ComfyUI using &quot;Export (API)&quot; and paste the JSON here.
</p>
</div>
</div>

View file

@ -28,9 +28,9 @@ export interface LLMConfig {
PEXELS_API_KEY?: string;
PIXABAY_API_KEY?: string;
// Local Image Generation (ComfyUI)
LOCAL_IMAGE_URL?: string;
LOCAL_IMAGE_WORKFLOW?: string; // ComfyUI workflow JSON
// ComfyUI
COMFYUI_URL?: string;
COMFYUI_WORKFLOW?: string;
// Other Configs
TOOL_CALLS?: boolean;

View file

@ -61,13 +61,13 @@ export const IMAGE_PROVIDERS: Record<string, ImageProviderOption> = {
apiKeyField: "GOOGLE_API_KEY",
apiKeyFieldLabel: "Google API Key"
},
local: {
value: "local",
comfyui: {
value: "comfyui",
label: "ComfyUI",
description: "Use your local ComfyUI server with custom workflows",
icon: "/icons/local.png",
icon: "/icons/comfyui.png",
requiresApiKey: false,
apiKeyField: "LOCAL_IMAGE_URL",
apiKeyField: "COMFYUI_URL",
apiKeyFieldLabel: "ComfyUI Server URL"
},
};

View file

@ -48,8 +48,8 @@ export const updateLLMConfig = (
disable_thinking: "DISABLE_THINKING",
extended_reasoning: "EXTENDED_REASONING",
web_grounding: "WEB_GROUNDING",
local_image_url: "LOCAL_IMAGE_URL",
local_image_workflow: "LOCAL_IMAGE_WORKFLOW",
comfyui_url: "COMFYUI_URL",
comfyui_workflow: "COMFYUI_WORKFLOW",
};
const configKey = fieldMappings[field];

View file

@ -73,8 +73,8 @@ export const hasValidLLMConfig = (llmConfig: LLMConfig) => {
return llmConfig.OPENAI_API_KEY && llmConfig.OPENAI_API_KEY !== "";
case "gemini_flash":
return llmConfig.GOOGLE_API_KEY && llmConfig.GOOGLE_API_KEY !== "";
case "local":
return llmConfig.LOCAL_IMAGE_URL && llmConfig.LOCAL_IMAGE_URL !== "";
case "comfyui":
return llmConfig.COMFYUI_URL && llmConfig.COMFYUI_URL !== "";
default:
return false;
}

View file

@ -96,8 +96,8 @@ const setupUserConfigFromEnv = () => {
process.env.EXTENDED_REASONING || existingConfig.EXTENDED_REASONING,
WEB_GROUNDING: process.env.WEB_GROUNDING || existingConfig.WEB_GROUNDING,
USE_CUSTOM_URL: process.env.USE_CUSTOM_URL || existingConfig.USE_CUSTOM_URL,
LOCAL_IMAGE_URL: process.env.LOCAL_IMAGE_URL || existingConfig.LOCAL_IMAGE_URL,
LOCAL_IMAGE_WORKFLOW: process.env.LOCAL_IMAGE_WORKFLOW || existingConfig.LOCAL_IMAGE_WORKFLOW,
COMFYUI_URL: process.env.COMFYUI_URL || existingConfig.COMFYUI_URL,
COMFYUI_WORKFLOW: process.env.COMFYUI_WORKFLOW || existingConfig.COMFYUI_WORKFLOW,
};
writeFileSync(userConfigPath, JSON.stringify(userConfig));