No output when using guardrails "content safety check output" with a custom langchain_core.Runnable

[abhijeetsourav]

I've made a custom langchain_core.Runnable to get model output of a diffusion pipeline, converting the output to a base64encoded url, to be used for content safety check and other output checks by the guardrails with vision-language model defined in config.yml. The class created is shown below:

from langchain_core.runnables import Runnable
from diffusers import DiffusionPipeline
import langchain_core
from nemoguardrails import LLMRails, RailsConfig
from langchain_core.prompts import ChatPromptTemplate

import torch
import base64
from io import BytesIO
from PIL import Image



class SDXLPipelineRunnable(Runnable):
    def __init__(self, base_pipe: DiffusionPipeline, refiner_pipe: DiffusionPipeline, rails_config: RailsConfig, n_steps=40, high_noise_frac=0.8):
        self.base = base_pipe
        self.refiner = refiner_pipe
        self.n_steps = n_steps
        self.high_noise_frac = high_noise_frac
        self.rails_config = rails_config
        self.rails = LLMRails(self.rails_config)
        self.buffer = BytesIO()
        self.model_output = ChatPromptTemplate.from_messages([
                {"role": "assistant", "content": [
                    {"type": "image_url", "image_url": {"url": "{img_url}"}},
                ]}
            ])
        

    def invoke(self, prompt, config=None, **kwargs):
        prompt = prompt.to_messages()[0].content

        latents = self.base(prompt=prompt, num_inference_steps=self.n_steps, denoising_end=self.high_noise_frac, output_type="latent").images
        image = self.refiner(prompt=prompt, num_inference_steps=self.n_steps, denoising_start=self.high_noise_frac, image=latents).images[0]
        print(image.size)      # image_size
        image = image.resize((320, 320), Image.LANCZOS)     # resize image
        image.save(self.buffer, format="PNG")
        img_bytes = self.buffer.getvalue()
        b64 = base64.b64encode(img_bytes).decode('utf-8')
        

        messages = [{
            "role" : 'assistant',
            "content" : [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{b64}",
                    }
                },
            ]}
        ]


        return messages

The SDXLPipelineRunnable along with prompt runnable is wrapped with nemoguardrails' RunnableRails as shown below:

import nest_asyncio
nest_asyncio.apply()

from nemoguardrails import RailsConfig
from nemoguardrails.integrations.langchain.runnable_rails import RunnableRails

config = RailsConfig.from_path("./config")
guardrails = RunnableRails(config)

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser, StrOutputParser

prompt = ChatPromptTemplate.from_messages([
    ("user", "{input}"),
])

output_parser = JsonOutputParser()

sdxl_runnable = SDXLPipelineRunnable(base, refiner, config, n_steps=10, high_noise_frac=0.8)
chain = prompt | sdxl_runnable 

guard_chain = guardrails | chain

res = guard_chain.invoke({"input": "some prompt..."})

the above all code gives the following result:
{'output': None}
that's it

Things I tried:

• I thought the output of my SDXLPipelineRunnable could be wrong, so i tried wrapping the messages in langchain_core.messages.AIMessage, but this also didn't work
• Then I took the output of chain.invoke() without guardrails to check how llm chain returns the output, I took the base64encoded string and passed it through a custom prompt to the same vision language model, I found out that the token limit of the message is reached, so I resized the pil image output from the difussion pipeline, but this also didn't help
  eventhough the max token limit error solved, guardrails showed no output

Important

I think the output from the custom runnable is not correct, the ouput format to be passed as bot_message to the prompt of content_safety_check_output is unknown to me, anyone who knows about this situation please help
Also , if any other method is there to do this please explain
I have also checked the Runnable-as-action section in the docs, but i don't know whether this wiil achieve me my desired results or not

Thank you very much!!!