refac: images

backend/open_webui/utils/files.py

@@ -1,5 +1,5 @@
 from open_webui.routers.images import (
-    load_b64_image_data,
+    get_image_data,
     upload_image,
 )
 
@@ -22,7 +22,7 @@ def get_image_url_from_base64(request, base64_image_string, metadata, user):
     if "data:image/png;base64" in base64_image_string:
         image_url = ""
         # Extract base64 image data from the line
-        image_data, content_type = load_b64_image_data(base64_image_string)
+        image_data, content_type = get_image_data(base64_image_string)
         if image_data is not None:
             image_url = upload_image(
                 request,

backend/open_webui/utils/headers.py

@@ -0,0 +1,11 @@
+from urllib.parse import quote
+
+
+def include_user_info_headers(headers, user):
+    return {
+        **headers,
+        "X-OpenWebUI-User-Name": quote(user.name, safe=" "),
+        "X-OpenWebUI-User-Id": user.id,
+        "X-OpenWebUI-User-Email": user.email,
+        "X-OpenWebUI-User-Role": user.role,
+    }

backend/open_webui/utils/images/comfyui.py

@@ -103,7 +103,7 @@ class ComfyUIWorkflow(BaseModel):
     nodes: list[ComfyUINodeInput]
 
 
-class ComfyUIGenerateImageForm(BaseModel):
+class ComfyUICreateImageForm(BaseModel):
     workflow: ComfyUIWorkflow
 
     prompt: str
@@ -116,8 +116,8 @@ class ComfyUIGenerateImageForm(BaseModel):
     seed: Optional[int] = None
 
 
-async def comfyui_generate_image(
-    model: str, payload: ComfyUIGenerateImageForm, client_id, base_url, api_key
+async def comfyui_create_image(
+    model: str, payload: ComfyUICreateImageForm, client_id, base_url, api_key
 ):
     ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://")
     workflow = json.loads(payload.workflow.workflow)

backend/open_webui/utils/middleware.py

@@ -45,9 +45,8 @@ from open_webui.routers.retrieval import (
     SearchForm,
 )
 from open_webui.routers.images import (
-    load_b64_image_data,
     image_generations,
-    GenerateImageForm,
+    CreateImageForm,
     upload_image,
 )
 from open_webui.routers.pipelines import (
@@ -770,7 +769,7 @@ async def chat_image_generation_handler(
     try:
         images = await image_generations(
             request=request,
-            form_data=GenerateImageForm(**{"prompt": prompt}),
+            form_data=CreateImageForm(**{"prompt": prompt}),
             user=user,
         )