Upload 5 files

README.md

@@ -1,10 +1,33 @@
----
-title: Image Gen
-emoji: ⚡
-colorFrom: indigo
-colorTo: pink
-sdk: docker
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Title: Image Generation with Chainlit
+Tags: [image, stability, langchain]
+
+# Image Generation with Chainlit
+
+This folder is showing how to use Stability AI to generate images and send them to the Chainlit UI.
+
+You will learn on how to use the [Image element](https://docs.chainlit.io/api-reference/elements/image) and integrate it with LangChain.
+
+## Description
+
+The provided code integrates Stability AI's image generation capabilities with Chainlit, a framework for building interactive web apps with Python. It includes functions to generate new images from text prompts and to edit existing images using prompts. The image generation is powered by the Stability AI's API, and the images are displayed using Chainlit's Image element.
+
+## Quickstart
+
+1. Ensure you have Chainlit installed and set up.
+2. Place your Stability AI API key in the environment variable `STABILITY_KEY`.
+3. Run `app.py` to start the Chainlit app.
+4. Use the `generate_image` function to create a new image from a text prompt.
+5. Use the `edit_image` function to edit an existing image with a new prompt.
+
+## Functions
+
+- `generate_image(prompt: str)`: Generates an image from a text prompt and returns the image name.
+- `edit_image(init_image_name: str, prompt: str)`: Edits an existing image based on the provided prompt and returns the new image name.
+
+## Tools
+
+- `generate_image_tool`: A Chainlit tool for generating images from text prompts.
+- `edit_image_tool`: A Chainlit tool for editing images with text prompts.
+
+To see the tools in action and interact with the generated images, follow the instructions in the [main readme](/README.md).
+![Rendering](./rendering.png)

app.py

@@ -0,0 +1,142 @@
+from langchain.agents import AgentExecutor, AgentType, initialize_agent
+from langchain.agents.structured_chat.prompt import SUFFIX
+from langchain.chat_models import ChatOpenAI
+from langchain.memory import ConversationBufferMemory
+from tools import edit_image_tool, generate_image_tool
+
+import chainlit as cl
+from chainlit.action import Action
+from chainlit.input_widget import Select, Switch, Slider
+
+
+@cl.action_callback("Create variation")
+async def create_variant(action: Action):
+    agent_input = f"Create a variation of {action.value}"
+    await cl.Message(content=f"Creating a variation of `{action.value}`.").send()
+    await main(cl.Message(content=agent_input))
+
+
+@cl.author_rename
+def rename(orig_author):
+    mapping = {
+        "LLMChain": "Assistant",
+    }
+    return mapping.get(orig_author, orig_author)
+
+
+@cl.cache
+def get_memory():
+    return ConversationBufferMemory(memory_key="chat_history")
+
+
+@cl.on_chat_start
+async def start():
+    settings = await cl.ChatSettings(
+        [
+            Select(
+                id="Model",
+                label="OpenAI - Model",
+                values=["gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-4", "gpt-4-32k"],
+                initial_index=1,
+            ),
+            Switch(id="Streaming", label="OpenAI - Stream Tokens", initial=True),
+            Slider(
+                id="Temperature",
+                label="OpenAI - Temperature",
+                initial=0,
+                min=0,
+                max=2,
+                step=0.1,
+            ),
+            Slider(
+                id="SAI_Steps",
+                label="Stability AI - Steps",
+                initial=30,
+                min=10,
+                max=150,
+                step=1,
+                description="Amount of inference steps performed on image generation.",
+            ),
+            Slider(
+                id="SAI_Cfg_Scale",
+                label="Stability AI - Cfg_Scale",
+                initial=7,
+                min=1,
+                max=35,
+                step=0.1,
+                description="Influences how strongly your generation is guided to match your prompt.",
+            ),
+            Slider(
+                id="SAI_Width",
+                label="Stability AI - Image Width",
+                initial=512,
+                min=256,
+                max=2048,
+                step=64,
+                tooltip="Measured in pixels",
+            ),
+            Slider(
+                id="SAI_Height",
+                label="Stability AI - Image Height",
+                initial=512,
+                min=256,
+                max=2048,
+                step=64,
+                tooltip="Measured in pixels",
+            ),
+        ]
+    ).send()
+    await setup_agent(settings)
+
+
+@cl.on_settings_update
+async def setup_agent(settings):
+    print("Setup agent with following settings: ", settings)
+
+    llm = ChatOpenAI(
+        temperature=settings["Temperature"],
+        streaming=settings["Streaming"],
+        model=settings["Model"],
+    )
+    memory = get_memory()
+    _SUFFIX = "Chat history:\n{chat_history}\n\n" + SUFFIX
+
+    agent = initialize_agent(
+        llm=llm,
+        tools=[generate_image_tool, edit_image_tool],
+        agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+        memory=memory,
+        agent_kwargs={
+            "suffix": _SUFFIX,
+            "input_variables": ["input", "agent_scratchpad", "chat_history"],
+        },
+    )
+    cl.user_session.set("agent", agent)
+
+
+@cl.on_message
+async def main(message: cl.Message):
+    agent = cl.user_session.get("agent")  # type: AgentExecutor
+    cl.user_session.set("generated_image", None)
+
+    # No async implementation in the Stability AI client, fallback to sync
+    res = await cl.make_async(agent.run)(
+        input=message.content, callbacks=[cl.LangchainCallbackHandler()]
+    )
+
+    elements = []
+    actions = []
+
+    generated_image_name = cl.user_session.get("generated_image")
+    generated_image = cl.user_session.get(generated_image_name)
+    if generated_image:
+        elements = [
+            cl.Image(
+                content=generated_image,
+                name=generated_image_name,
+                display="inline",
+            )
+        ]
+        actions = [cl.Action(name="Create variation", value=generated_image_name)]
+
+    await cl.Message(content=res, elements=elements, actions=actions).send()

chainlit.md

@@ -0,0 +1,24 @@
+# 🎨 Chainlit Image Gen demo
+
+Welcome to our creative image generation demo built with [Chainlit](https://chainlit.io), [LangChain](https://python.langchain.com/en/latest/index.html), and [Stability AI](https://stability.ai/)! 🌟 This app allows you to create and edit unique images simply by chatting with it. Talk about having an artistic conversation! 🎨🗨️
+
+This demo has also been adapted to use the new [ChatSettings](https://docs.chainlit.io/concepts/chat-settings) feature introduced in chainlit `0.6.2`. You are now able to tweak Stability AI settings to your liking!
+
+## 🎯 Example
+
+Try asking:
+```
+Anime style snowy swiss mountains lit by the moon. Clear night sky. HD.
+```
+
+![Result](https://chainlit-cloud.s3.eu-west-3.amazonaws.com/docs/img-gen-0.jpg)
+
+You can then ask for modifications:
+```
+change the clear night sky with a starry sky
+```
+
+![Result](https://chainlit-cloud.s3.eu-west-3.amazonaws.com/docs/img-gen-1.jpg)
+
+## ⚠️ Disclaimer
+Please note that the primary goal of this demo is to showcase the ease and convenience of building LLM apps using Chainlit and other tools rather than presenting a state-of-the-art image generation application. 

tools.py

@@ -0,0 +1,108 @@
+import io
+import os
+
+import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
+from langchain.tools import StructuredTool, Tool
+from PIL import Image
+from stability_sdk import client
+
+import chainlit as cl
+
+os.environ["STABILITY_HOST"] = "grpc.stability.ai:443"
+
+
+def get_image_name():
+    image_count = cl.user_session.get("image_count")
+    if image_count is None:
+        image_count = 0
+    else:
+        image_count += 1
+
+    cl.user_session.set("image_count", image_count)
+
+    return f"image-{image_count}"
+
+
+def _generate_image(prompt: str, init_image=None):
+    # Set up our connection to the API.
+    stability_api = client.StabilityInference(
+        key=os.environ["STABILITY_KEY"],  # API Key reference.
+        verbose=True,  # Print debug messages.
+        engine="stable-diffusion-xl-beta-v2-2-2",  # Set the engine to use for generation.
+        # Available engines: stable-diffusion-v1 stable-diffusion-v1-5 stable-diffusion-512-v2-0 stable-diffusion-768-v2-0
+        # stable-diffusion-512-v2-1 stable-diffusion-768-v2-1 stable-diffusion-xl-beta-v2-2-2 stable-inpainting-v1-0 stable-inpainting-512-v2-0
+    )
+
+    start_schedule = 0.8 if init_image else 1
+
+    cl_chat_settings = cl.user_session.get("chat_settings")
+
+    # Set up our initial generation parameters.
+    answers = stability_api.generate(
+        prompt=prompt,
+        init_image=init_image,
+        start_schedule=start_schedule,
+        seed=992446758,  # If a seed is provided, the resulting generated image will be deterministic.
+        # What this means is that as long as all generation parameters remain the same, you can always recall the same image simply by generating it again.
+        # Note: This isn't quite the case for CLIP Guided generations, which we tackle in the CLIP Guidance documentation.
+        steps=int(cl_chat_settings["SAI_Steps"]),  # Amount of inference steps performed on image generation. Defaults to 30.
+        cfg_scale=cl_chat_settings["SAI_Cfg_Scale"],  # Influences how strongly your generation is guided to match your prompt.
+        # Setting this value higher increases the strength in which it tries to match your prompt.
+        # Defaults to 7.0 if not specified.
+        width=int(cl_chat_settings["SAI_Width"]),  # Generation width, defaults to 512 if not included.
+        height=int(cl_chat_settings["SAI_Height"]),  # Generation height, defaults to 512 if not included.
+        samples=1,  # Number of images to generate, defaults to 1 if not included.
+        sampler=generation.SAMPLER_K_EULER  # Choose which sampler we want to denoise our generation with.
+        # Defaults to k_dpmpp_2m if not specified. Clip Guidance only supports ancestral samplers.
+        # (Available Samplers: ddim, plms, k_euler, k_euler_ancestral, k_heun, k_dpm_2, k_dpm_2_ancestral, k_dpmpp_2s_ancestral, k_lms, k_dpmpp_2m, k_dpmpp_sde)
+    )
+
+    # Set up our warning to print to the console if the adult content classifier is tripped.
+    # If adult content classifier is not tripped, save generated images.
+    for resp in answers:
+        for artifact in resp.artifacts:
+            if artifact.finish_reason == generation.FILTER:
+                raise ValueError(
+                    "Your request activated the API's safety filters and could not be processed."
+                    "Please modify the prompt and try again."
+                )
+            if artifact.type == generation.ARTIFACT_IMAGE:
+                name = get_image_name()
+                cl.user_session.set(name, artifact.binary)
+                cl.user_session.set("generated_image", name)
+                return name
+            else:
+                raise ValueError(
+                    f"Your request did not generate an image. Please modify the prompt and try again. Finish reason: {artifact.finish_reason}"
+                )
+
+
+def generate_image(prompt: str):
+    image_name = _generate_image(prompt)
+    return f"Here is {image_name}."
+
+
+def edit_image(init_image_name: str, prompt: str):
+    init_image_bytes = cl.user_session.get(init_image_name)
+    if init_image_bytes is None:
+        raise ValueError(f"Could not find image `{init_image_name}`.")
+
+    init_image = Image.open(io.BytesIO(init_image_bytes))
+    image_name = _generate_image(prompt, init_image)
+
+    return f"Here is {image_name} based on {init_image_name}."
+
+
+generate_image_tool = Tool.from_function(
+    func=generate_image,
+    name="GenerateImage",
+    description="Useful to create an image from a text prompt.",
+    return_direct=True,
+)
+
+edit_image_tool = StructuredTool.from_function(
+    func=edit_image,
+    name="EditImage",
+    description="Useful to edit an image with a prompt. Works well with commands such as 'replace', 'add', 'change', 'remove'.",
+    return_direct=True,
+)