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.