returning an image to open-webui via a pipeline

[mrazmat]

hi all,

I have a a pipeline which loads and functions ok. It connects to replicate and sends my prompt where an image is generated. The output of the replicate.run () command is an list of outputs including a link to the image that is generated. I want to send this back to open-webui to display. however, i am unable to. can someone help me?

I think my issue is with the pipe function and the return type of pipe and open-webui is rendering the return from my pipeline.

Any assistance would be very much appreciated!

import base64
import os
import replicate
import logging
import requests
from PIL import Image
from urllib.request import urlretrieve
from typing import List, Union, Generator, Iterator
#from pydantic import BaseModel
from io import BytesIO
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Pipeline:
    def __init__(self):
        self.name = "Replicate Image Generator"
        self.api_token = "my key"
        if not self.api_token:
            raise ValueError("REPLICATE_API_TOKEN environment variable not set.")
        self.client = replicate.Client(api_token=self.api_token)
        self.model =  "black-forest-labs/flux-schnell:f2ab8a5bfe79f02f0789a146cf5e73d2a4ff2684a98c2b303d1e1ff3814271db"
        self.API_URL = "https://api.replicate.com/v1/predictions"

        pass

    async def on_startup(self):
        # This function is called when the server is started.
        print(f"on_startup:{__name__}")
        pass 

    async def on_shutdown(self):
        # This function is called when the server is shutdown.
        print(f"on_shutdown:{__name__}")
        pass

    def pipe(self, user_message: str, model_id: str, messages: List[dict], body: dict) -> Union[Iterator[bytes], bytes]:
        # Generate or retrieve the image
        image_url_or_data = self._get_or_generate_image(user_message)

        if isinstance(image_url_or_data, str):
            from urllib.request import urlopen
            img = Image.open(urlopen(image_url_or_data))
        else:
            img = Image.open(BytesIO(image_url_or_data))
        
        # Convert the image to bytes
        buffered = BytesIO()
        img.save(buffered, format="PNG")
        img_bytes = buffered.getvalue()
        return img_bytes

    def _get_or_generate_image(self, prompt: str):
        output = self.client.run( 
                 self.model,
                 input={
                     "prompt": prompt,
                     "num_outputs": 1,
                     "aspect_ratio": "1:1",
                     "output_format": "png",
                     "output_quality": 90
                 }
        )
        return output[0]

[Akatsuki-Ryu]

take a look at this example , i was wondering the same question . dbb6800

[Akatsuki-Ryu]

if you are looking for flux integration . i made this
https://github.com/Akatsuki-Ryu/openwebui-pipelines/blob/main/examples/pipelines/integrations/replicate_flux_pipeline.py

[mrazmat]

Thanks! Your pipeline worked. I was able to send my prompt through the pipeline to FLUX on replicate. It generated an image, however on the response back to openwebui it refused to display. I tried the various output formats, but it still doesn't want to display on the UI.

[Akatsuki-Ryu]

sorry about the late reply. i think they made some changes on the API. please checkout my latest code

[mrazmat]

Thanks. Odd I am unable to see your latest replicate_flux_pipeline.py. This looks like it was last edited in October. I'm likely missing something :|

https://github.com/Akatsuki-Ryu/openwebui-pipelines/blob/main/examples/pipelines/integrations/replicate_flux_pipeline.py

Can you point me in the right direction?

[Akatsuki-Ryu]

there are some changes again lately. let me find you the latest file , I will come back to you

[Akatsuki-Ryu]

i think this is the latest code that i have , please have a try , if problem , please let me know

[Hsun1128]

I used return image_data successfully and can see the image on the open webui. However, when I try to use itertools.chain to combine r.iter_lines() and image_iter, I can only see the content, but not the image. Does anyone know what is wrong?

This is just a test. I hope to attach the image after the response is sent, but I haven’t passed the image to the model for processing yet.

from typing import List, Union, Generator, Iterator
from schemas import OpenAIChatMessage
import requests
import json
import base64
from io import BytesIO
from PIL import Image
import itertools



class Pipeline:
    def __init__(self):
        # Optionally, you can set the id and name of the pipeline.
        # Best practice is to not specify the id so that it can be automatically inferred from the filename, so that users can install multiple versions of the same pipeline.
        # The identifier must be unique across all pipelines.
        # The identifier must be an alphanumeric string that can include underscores or hyphens. It cannot contain spaces, special characters, slashes, or backslashes.
        # self.id = "ollama_pipeline"
        self.name = "Ollama Pipeline"
        pass

    async def on_startup(self):
        # This function is called when the server is started.
        print(f"on_startup:{__name__}")
        pass

    async def on_shutdown(self):
        # This function is called when the server is stopped.
        print(f"on_shutdown:{__name__}")
        pass

    def process_messages(self, messages):
        # Collection photo on last message
        try:
            image_url = messages[-1]['content'][1]['image_url']['url']
            
            image_data = f"\n![image]({image_url})\n"
        # cut Base64 tag
#        if image_data.startswith("data:image"):
#            image_data = image_data.split(",")[1]

        # Base64 to 2進制
#        image_binary = base64.b64decode(image_data)

        # transfer to PIL image
#        image = Image.open(BytesIO(image_binary))
        except:
            image_data = 'None image in user_message'

        return image_data


    def pipe(
        self, user_message: str, model_id: str, messages: List[dict], body: dict
    ) -> Union[str, Generator, Iterator]:
        # This is where you can add your custom pipelines like RAG.
        print(f"pipe:{__name__}")

        OLLAMA_BASE_URL = "http://host.docker.internal:11434"
        MODEL = "llama3"

        if "user" in body:
            print("######################################")
            print(f'# User: {body["user"]["name"]} ({body["user"]["id"]})')
            print(f"# Message: {user_message}")
            print("######################################")

        try:
            r = requests.post(
                url=f"{OLLAMA_BASE_URL}/v1/chat/completions",
                json={**body, "model": MODEL},
                stream=True,
            )

            r.raise_for_status()

            if False:#body["stream"]:
                return r.iter_lines()
            else:
                image_data = self.process_messages(messages)
                image_iter = iter([image_data])
                combined = itertools.chain(r.iter_lines(), image_iter)
                return combined #r.json()
        except Exception as e:
            return f"Error: {e}"