Toggle Menu
Prev Up Next

# Source code for langchain_nvidia_ai_endpoints.image_gen

    """Embeddings Components Derived from NVEModel/Embeddings"""
    import base64
    from io import BytesIO
    from typing import Any, List, Optional

    import requests
    from langchain_core.callbacks.manager import CallbackManagerForLLMRun
    from langchain_core.language_models import LLM
    from langchain_core.pydantic_v1 import Field
    from langchain_core.runnables import Runnable, RunnableLambda
    from PIL import Image

    from langchain_nvidia_ai_endpoints._common import _NVIDIAClient

    """
    ## Image Generation Models

    Due to the similarity of the underlying API, a selection of **Image Generation Models**
    can be supported using the LLM interface. One example is the [**Stable Diffusion XL**]
    (https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/sdxl/api) model,
    which expects a prompt and some other arguments as input and produces an image (passed
    back as a b64-encoded string).

    ```
    from langchain_nvidia_ai_endpoints import ImageGenNVIDIA

    ImageGenNVIDIA.get_available_models()  ## Only shows models supported by ImageGenNVIDIA

    sdxl = ImageGenNVIDIA(model="sdxl")

    out = sdxl.invoke("A picture of a cat, aesthetic")
    out[:100] + "..."
    ```

    To see the image, we can either convert it manually or use the built-in conversion
    utility

    ```
    from langchain_nvidia_ai_endpoints.image_gen import ImageParser

    # from io import BytesIO
    # from PIL import Image
    # import base64

    # Image.open(BytesIO(base64.decodebytes(bytes(out[10:], "utf-8"))))
    ImageParser().invoke(out)  ## Runnable that does it all for you
    ```

    In addition to the prompt, we can do a bit of hyperparameter tweaking and add some
    negative prompt components that we'd like to avoid. We will also use the `ImageParser`
    runnable automatically by calling the `.as_pil()` method.

    ```
    sdxl = ImageGenNVIDIA(
        inference_steps = 100,
        negative_prompt = "ugly,bad eyes,low-res,crooked nose, smudged, painted",
    )

    sdxl.as_pil().invoke("A picture of a big green dog, futuristic cyberpunk")
    ```

    Note that under the hood, `as_pil` returns a merger of the `ImageGenModel` object with
    the `ImageParser` output parser. As a result, you may have trouble interacting with the
    aggregation. Note that you can reference the first half of the pipeline via `.first` or
    similar.

    ## Example of image generation with OpenAI and DALL-E

    ```
    from langchain_nvidia_ai_endpoints import ImageGenNVIDIA
    from getpass import getpass
    import os

    if not os.environ.get("OPENAI_API_KEY", "").startswith("sk-"):
        os.environ["OPENAI_API_KEY"] = getpass("Enter your OPENAI Key: ")

    llm = ImageGenNVIDIA().mode("openai")
    llm.available_models
    ```

    ```
    from langchain_nvidia_ai_endpoints import ImageGenNVIDIA

    dalle = ImageGenNVIDIA().mode("openai", model="dall-e-3")

    def payload_fn(d):
        if d:
            drop_keys = ["guidance_scale", "seed", "negative_prompt", "sampler"]
            d = {k: v for k, v in d.items() if k not in drop_keys}
        return d

    dalle.client.payload_fn = payload_fn

    dalle.as_pil().invoke("City skyline, neon green tint, aesthetic futuristic realistic")
    print(dalle.client.last_inputs['json'])
    dalle.client.last_response.json()
    ```
    """

    def _get_pil_from_response(data: str) -> Image.Image:
        if data.startswith("url: "):
            body = requests.get(data[4:], stream=True).raw
        elif data.startswith("b64_json: "):
            body = BytesIO(base64.decodebytes(bytes(data[10:], "utf-8")))
        else:
            raise ValueError(f"Invalid response format: {str(data)[:100]}")
        return Image.open(body)

    def ImageParser() -> RunnableLambda[str, Image.Image]:
        return RunnableLambda(_get_pil_from_response)

    class ImageGenNVIDIA(_NVIDIAClient, LLM):
        """NVIDIA's AI Foundation Retriever Question-Answering Asymmetric Model."""

        _default_model: str = "sdxl"
        infer_endpoint: str = Field("{base_url}/images/generations")
        model: str = Field(_default_model, description="Name of the model to invoke")
        negative_prompt: Optional[str] = Field(description="Sampling temperature in [0, 1]")
        sampler: Optional[str] = Field(description="Sampling strategy for process")
        guidance_scale: Optional[float] = Field(description="The scale of guidance")
        seed: Optional[int] = Field(description="The seed for deterministic results")

        @property
        def _llm_type(self) -> str:
            """Return type of NVIDIA AI Foundation Model Interface."""
            return "nvidia-image-gen-model"

        def _call(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> str:
            """Run the Image Gen Model on the given prompt and input."""
            payload = {
                "prompt": prompt,
                "negative_prompt": kwargs.get("negative_prompt", self.negative_prompt),
                "sampler": kwargs.get("sampler", self.sampler),
                "guidance_scale": kwargs.get("guidance_scale", self.guidance_scale),
                "seed": kwargs.get("seed", self.seed),
            }
            if self.get_binding_model():
                payload["model"] = self.get_binding_model()
            response = self.client.get_req(
                model_name=self.model, payload=payload, endpoint="infer"
            )
            response.raise_for_status()
            out_dict = response.json()
            if "data" in out_dict:
                out_dict = out_dict.get("data")[0]
            if "url" in out_dict:
                output = "url: {}".format(out_dict.get("url"))
            elif "b64_json" in out_dict:
                output = "b64_json: {}".format(out_dict.get("b64_json"))
            else:
                output = str(out_dict)
            return output

        def as_pil(self, **kwargs: Any) -> Runnable:
            """Returns a model that outputs a PIL image by default"""
            return self | ImageParser(**kwargs)