import copy
import logging
import os
import os.path as osp
from functools import partial
from pydoc import locate
import shutil
import json
from traceback import print_exc
import uuid
from pathlib import Path
from collections import OrderedDict
import numpy as np
from PIL import Image
import whisper
import fire
import gradio as gr
import gradio.themes.base as ThemeBase
from gradio.themes.utils import colors, sizes
from gradio.components.image_editor import Brush
import sys
sys.path.append(os.getcwd())
from cllm.agents.builtin import plans
from cllm.services.general.api import remote_logging
from cllm.agents import container, FILE_EXT
from cllm.utils import get_real_path, plain2md, md2plain
import openai
openai.api_base = os.environ.get("OPENAI_API_BASE", None)
openai.api_key = os.environ.get("OPENAI_API_KEY", None)
logging.basicConfig(
filename="cllm.log",
level=logging.INFO,
format="%(asctime)s %(levelname)-8s %(message)s",
)
logger = logging.getLogger(__name__)
RESOURCE_ROOT = os.environ.get("CLIENT_ROOT", "./client_resources")
def is_image(file_path):
ext = FILE_EXT["image"]
_, extension = os.path.splitext(file_path)
return extension[1:] in ext
def is_video(file_path):
ext = FILE_EXT["video"]
_, extension = os.path.splitext(file_path)
return extension[1:] in ext
def is_audio(file_path):
ext = FILE_EXT["audio"]
_, extension = os.path.splitext(file_path)
return extension[1:] in ext
def get_file_type(file_path):
if is_image(file_path):
if "mask" in file_path:
return "mask"
return "image"
elif is_video(file_path):
return "video"
elif is_audio(file_path):
return "audio"
raise ValueError("Invalid file type")
def convert_dict_to_frame(data):
import pandas
outputs = []
for k, v in data.items():
output = {"Resource": k}
if not isinstance(v, str):
output["Type"] = str(v.__class__)
else:
output["Type"] = v
outputs.append(output)
if len(outputs) == 0:
return None
return pandas.DataFrame(outputs)
class Seafoam(ThemeBase.Base):
def __init__(
self,
*,
primary_hue=colors.emerald,
secondary_hue=colors.blue,
neutral_hue=colors.gray,
spacing_size=sizes.spacing_md,
radius_size=sizes.radius_md,
text_size=sizes.text_sm,
):
super().__init__(
primary_hue=primary_hue,
secondary_hue=secondary_hue,
neutral_hue=neutral_hue,
spacing_size=spacing_size,
radius_size=radius_size,
text_size=text_size,
)
super().set(
body_background_fill_dark="#111111",
button_primary_background_fill="*primary_300",
button_primary_background_fill_hover="*primary_200",
button_primary_text_color="black",
button_secondary_background_fill="*secondary_300",
button_secondary_background_fill_hover="*secondary_200",
border_color_primary="#0BB9BF",
slider_color="*secondary_300",
slider_color_dark="*secondary_600",
block_title_text_weight="600",
block_border_width="3px",
block_shadow="*shadow_drop_lg",
button_shadow="*shadow_drop_lg",
button_large_padding="10px",
)
class InteractionLoop:
def __init__(
self,
controller="cllm.agents.code.Controller",
):
self.stream = True
Controller = locate(controller)
self.controller = Controller(stream=self.stream, interpretor_kwargs=dict())
self.whisper = whisper.load_model("base")
def _gen_new_name(self, r_type, ext="png"):
this_new_uuid = str(uuid.uuid4())[:6]
new_file_name = f"{this_new_uuid}_{r_type}.{ext}"
return new_file_name
def init_state(self):
user_state = OrderedDict()
user_state["resources"] = OrderedDict()
user_state["history_msgs"] = []
resources = OrderedDict()
for item in sorted(os.listdir("./assets/resources")):
if item.startswith("."):
continue
shutil.copy(
osp.join("./assets/resources", item),
osp.join(RESOURCE_ROOT, item),
)
resources[item] = get_file_type(item)
# return user_state, user_state["resources"]
return user_state, resources
def add_file(self, user_state, history, file):
if user_state.get("resources", None) is None:
user_state["resources"] = OrderedDict()
if file is None:
return user_state, None, history, None
# filename = os.path.basename(file.name)
file = Path(file)
ext = file.suffix[1:]
if ext in FILE_EXT["image"]:
ext = "png"
r_type = get_file_type(file.name)
new_filename = self._gen_new_name(r_type, ext)
saved_path = get_real_path(new_filename)
if ext in FILE_EXT["image"]:
Image.open(file).convert("RGB").save(saved_path, "png")
user_state["input_image"] = new_filename
else:
shutil.copy(file, saved_path)
logger.info(f"add file: {saved_path}")
user_state["resources"][new_filename] = r_type
for key, val in user_state["resources"].items():
if key == "prompt_points":
user_state["resources"].pop(key)
break
history, _ = self.add_text(history, (saved_path,), role="human", append=False)
history, _ = self.add_text(
history, f"Recieved file {new_filename}", role="assistant", append=False
)
memory = convert_dict_to_frame(user_state["resources"])
image_name = None
if Path(saved_path).suffix[1:] in FILE_EXT["image"]:
image_name = saved_path
return user_state, image_name, history, memory
def add_msg(self, history, text, audio, role="assistant", append=False):
if text is not None and text.strip() != "":
return self.add_text(history, text, role=role, append=append)
elif audio is not None:
return self.add_audio(history, audio, role=role, append=append)
return history, ""
def add_sketch(self, user_state, history, sketch):
if user_state.get("resources", None) is None:
user_state["resources"] = OrderedDict()
if sketch is None or "layers" not in sketch:
return user_state, None, history, None
mask = None
for layer in sketch["layers"]:
alpha = layer[:, :, 3:] // 255
if mask is None:
mask = np.ones_like(layer[:, :, :3]) * 255
mask = mask * (1 - alpha) + layer[:, :, :3] * alpha
ext = "png"
r_type = "scribble"
new_filename = self._gen_new_name(r_type, ext)
saved_path = get_real_path(new_filename)
if ext in FILE_EXT["image"]:
Image.fromarray(mask).save(saved_path, "png")
user_state["sketch_image"] = new_filename
logger.info(f"add file: {saved_path}")
user_state["resources"][new_filename] = r_type
history, _ = self.add_text(history, (saved_path,), role="human", append=False)
history, _ = self.add_text(
history, f"Recieved file {new_filename}", role="assistant", append=False
)
memory = convert_dict_to_frame(user_state["resources"])
return user_state, history, memory
def add_text(self, history, text, role="assistant", append=False):
if history is None:
history = []
# return history, ""
assert role in ["human", "assistant"]
idx = 0
if len(history) == 0 or role == "human":
history.append([None, None])
if role == "assistant":
idx = 1
if not append and history[-1][1] is not None:
history.append([None, None])
if append:
history[-1][idx] = (
text if history[-1][idx] is None else history[-1][idx] + text
)
else:
history[-1][idx] = text
if isinstance(text, str):
logger.info(f"add text: {md2plain(text)}")
return history, ""
def add_audio(self, history, audio, role="assistant", append=False):
assert role in ["human", "assistant"]
result = self.whisper.transcribe(audio)
text = result["text"]
logger.info(f"add audio: {text}")
return self.add_text(history, text, role=role, append=append)
def plan(self, user_state, input_image, history, history_plan):
logger.info(f"Task plan...")
if user_state.get("resources", None) is None:
user_state["resources"] = OrderedDict()
request = history[-1][0]
user_state["request"] = request
if isinstance(request, str) and request.startswith("$"):
solution = f'show$("{request[1:]}")'
else:
solution = self.controller.plan(request, state=user_state)
print(f"request: {request}")
if solution == self.controller.SHORTCUT:
# md_text = "**Using builtin shortcut solution.**"
history, _ = self.add_text(
history, solution, role="assistant", append=False
)
user_state["solution"] = solution
user_state["history_msgs"] = history
yield user_state, input_image, history, [solution]
elif isinstance(solution, str) and solution.startswith("show$"):
user_state["solution"] = solution
yield user_state, input_image, history, solution
else:
output_text = (
"The whole process will take some time, please be patient.
"
)
history, _ = self.add_text(
history, output_text, role="assistant", append=True
)
yield user_state, input_image, history, history_plan
task_decomposition = next(solution)
if task_decomposition in [None, [], ""]:
output = "Error: unrecognized resource(s) in task decomposition."
task_decomposition = "[]"
else:
output = task_decomposition
output = f"**Task Decomposition:**\n{output}"
output = plain2md(output)
history, _ = self.add_text(history, output, role="assistant", append=True)
user_state["task_decomposition"] = json.loads(task_decomposition)
yield user_state, input_image, history, history_plan
history, _ = self.add_text(
history,
plain2md("\n\n**Thoughs-on-Graph:**\n"),
role="assistant",
append=True,
)
yield user_state, input_image, history, history_plan
solution_str = next(solution)
logger.info(f"Thoughs-on-Graph: \n{solution_str}")
if solution_str in [None, [], ""]:
output = "Empty solution possibly due to some internal errors."
solution_str = "[]"
else:
output = solution_str
output_md = plain2md(output)
history, _ = self.add_text(
history, output_md, role="assistant", append=True
)
solution = json.loads(solution_str)
user_state["solution"] = solution
user_state["history_msgs"] = history
yield user_state, input_image, history, solution
def execute(self, user_state, input_image, history, history_plan):
resources_state = user_state.get("resources", OrderedDict())
solution = user_state.get("solution", None)
if not solution:
yield user_state, input_image, history, history_plan
return
logger.info(f"Tool execution...")
if isinstance(solution, str) and solution.startswith("show$"):
key = solution[7:-2]
r_type = resources_state.get(key)
if r_type is None:
resource = f"{key} not found"
resource = container.auto_type("None", r_type, key)
history, _ = self.add_text(
history, (resource.to_chatbot(),), role="assistant"
)
user_state["history_msgs"] = history
yield user_state, input_image, history, history_plan
return
elif solution:
results = self.controller.execute(solution, state=user_state)
if not results:
yield user_state, input_image, history, history_plan
return
user_state["outputs"] = []
for result_per_step, executed_solutions, wrapped_outputs in results:
tool_name = json.dumps(result_per_step[0], ensure_ascii=False)
args = json.dumps(result_per_step[1], ensure_ascii=False)
ret = json.dumps(result_per_step[2], ensure_ascii=False)
history, _ = self.add_text(
history,
f"Call **{tool_name}:**
**Args**: {plain2md(args)}
**Ret**: {plain2md(ret)}",
role="assistant",
)
user_state["history_msgs"] = history
user_state["executed_solutions"] = executed_solutions
yield user_state, input_image, history, history_plan
for _, output in enumerate(wrapped_outputs):
if output is None or output.value is None:
continue
if isinstance(output, container.File):
history, _ = self.add_text(
history,
f"Here is {output.filename}:",
role="assistant",
)
history, _ = self.add_text(
history, (output.to_chatbot(),), role="assistant"
)
user_state["outputs"].extend(wrapped_outputs)
user_state["history_msgs"] = history
yield user_state, input_image, history, history_plan
else:
yield user_state, input_image, history, history_plan
def reply(self, user_state, history):
logger.info(f"Make response...")
executed_solution = user_state.get("executed_solutions", None)
resources_state = user_state.get("resources", OrderedDict())
solution = user_state.get("solution", None)
memory = convert_dict_to_frame(resources_state)
if isinstance(solution, str) and solution.startswith("show$"):
return user_state, history, memory
outputs = user_state.get("outputs", None)
response, user_state = self.controller.reply(
executed_solution, outputs, user_state
)
# prompt_mask_out = None
for i, output in enumerate(response):
if isinstance(output, container.File):
history, _ = self.add_text(history, f"Here is [{output.filename}]: ")
history, _ = self.add_text(history, (output.to_chatbot(),))
elif i == 0:
history, _ = self.add_text(history, output.to_chatbot())
user_state["history_msgs"] = history
return user_state, history, memory
def vote(self, user_state, history, data: gr.LikeData):
data_value = data.value
if isinstance(data_value, dict):
data_value = json.dumps(data_value)
if data.liked:
print("You upvoted this response: ", data_value)
logger.info("You upvoted this response: " + data_value)
else:
print("You downvoted this response: ", data_value)
logger.info("You downvoted this response: " + data_value)
remote_logging(
user_state.get("history_msgs", []),
user_state.get("task_decomposition", ""),
user_state.get("solution", []),
data_value,
data.liked,
)
msg = f"Thanks for your feedback! You feedback will contribute a lot to improving our ControlLLM."
history, _ = self.add_text(history, msg)
user_state["history_msgs"] = history
return user_state, history
def save_point(self, user_state, history, data: gr.SelectData):
if isinstance(data, gr.LikeData):
return self.vote(user_state, history, data)
if not isinstance(data, gr.SelectData):
return user_state, history
resource_state = user_state.get("resources")
input_image = user_state.get("input_image", None)
if input_image is None:
history, _ = self.add_text(history, "Please upload an image at first.")
history, _ = self.add_text(history, plans.BUILTIN_SEG_BY_POINTS, "human")
user_state["history_msg"] = history
return user_state, history
resource_state.pop(input_image, None)
resource_state[input_image] = "image"
history = history + [[plans.BUILTIN_SEG_BY_POINTS, None]]
points = []
if isinstance(points, str):
points = json.loads(points)
points.append(data.index)
resource_state[json.dumps(points)] = "prompt_points"
user_state["resources"] = resource_state
return user_state, history
def on_switch_input(state_input, text, audio, disable=False):
if state_input == "audio" or disable:
return "text", gr.update(visible=True), gr.update(visible=False)
return "audio", gr.update(visible=False), gr.update(visible=True)
def on_mask_submit(history):
history = history + [(plans.BUILTIN_SEG_BY_MASK, None)]
return history
def app(controller="cllm.agents.tog.Controller", https=False, **kwargs):
loop = InteractionLoop(controller=controller)
init_state, builtin_resources = loop.init_state()
css = """
code {
font-size: var(--text-sm);
white-space: pre-wrap; /* Since CSS 2.1 */
white-space: -moz-pre-wrap; /* Mozilla, since 1999 */
white-space: -pre-wrap; /* Opera 4-6 */
white-space: -o-pre-wrap; /* Opera 7 */
word-wrap: break-word; /* Internet Explorer 5.5+ */
}
"""
with gr.Blocks(theme=Seafoam(), css=css) as demo:
gr.HTML(
"""
ControlLLM
A framework for multi-modal interaction which is able to control LLMs over invoking tools more accurately.
GitHub
ArXiv
""",
)
state_input = gr.State("text")
user_state = gr.State(copy.deepcopy(init_state))
with gr.Row():
with gr.Column(scale=6):
with gr.Tabs():
with gr.Tab("Chat"):
chatbot = gr.Chatbot(
[],
elem_id="chatbot",
avatar_images=[
"assets/human.png",
"assets/assistant.png",
],
show_copy_button=True,
height=550,
)
with gr.Row():
with gr.Column(scale=12):
text = gr.Textbox(
show_label=False,
placeholder="Enter text and press enter, or upload an image.",
container=False,
)
audio = gr.Audio(
sources="microphone", type="filepath", visible=False
)
with gr.Column(scale=2, min_width=80):
submit = gr.Button("Submit", variant="primary")
with gr.Column(scale=1, min_width=40):
record = gr.Button("๐๏ธ")
with gr.Column(scale=1, min_width=40):
upload_btn = gr.UploadButton(
"๐",
file_types=[
"image",
"video",
"audio",
".pdf",
],
)
gr.Examples(
[
"Who are you?",
"How is about weather in Beijing",
"Describe the given image.",
"find the woman wearing the red skirt in the image",
"Generate a video that shows Pikachu surfing in waves.",
"How many horses are there in the image?",
"Can you erase the dog in the given image?",
"Remove the object based on the given mask.",
"Can you make a video of a serene lake with vibrant green grass and trees all around? And then create a webpage using HTML to showcase this video?",
"Generate an image that shows a beautiful landscape with a calm lake reflecting the blue sky and white clouds. Then generate a video to introduce this image.",
"replace the masked object with a cute yellow dog",
"replace the sheep with a cute dog in the image",
"Recognize the action in the video",
"Generate an image where a astronaut is riding a horse",
"Please generate a piece of music from the given image",
"Please give me an image that shows an astronaut riding a horse on mars.",
"Whatโs the weather situation in Berlin? Can you generate a new image that represents the weather in there?",
"Can you recognize the text from the image and tell me how much is Eggs Florentine?",
"Generate a piece of music for this video and dub this video with generated music",
"Generate a new image based on depth map from input image",
"Remove the cats from the image_1.png, image_2.png, image_3.png",
"I need the banana removed from the c4c40e_image.png, 9e867c_image.png, 9e13sc_image.png",
"I would be so happy if you could create a new image using the scribble from input image. The new image should be a tropical island with a dog. Write a detailed description of the given image. and highlight the dog in image",
"Please generate a piece of music and a new video from the input image",
"generate a new image conditioned on the segmentation from input image and the new image shows that a gorgeous lady is dancing",
"generate a new image with a different background but maintaining the same composition as input image",
"Generate a new image that shows an insect robot preparing a delicious meal. Then give me a video based on new image. Finally, dub the video with suitable background music.",
"Translate the text into speech: I have a dream that one day this nation will rise up and live out the true meaning of its creed: We hold these truths to be self-evident that all men are created equal.I have a dream that one day on the red hills of Georgia the sons of former slaves and the sons of former slave owners will be able to sit down together at the table of brotherhood. I have a dream that one day even the state of Mississippi, a state sweltering with the heat of injustice, sweltering with the heat of oppression, will be transformed into an oasis of freedom and justice.",
],
inputs=[text],
)
gr.Examples(
list(plans.BUILTIN_PLANS.keys()),
inputs=[text],
label="Builtin Examples",
)
with gr.Column(scale=5):
with gr.Tabs():
with gr.Tab("Click"):
click_image = gr.Image(
sources=["upload", "clipboard"],
interactive=True,
type="filepath",
)
with gr.Row():
click_image_submit_btn = gr.Button(
"Upload", variant="primary"
)
gr.Examples(
[
osp.join("./assets/resources", item)
for item in builtin_resources.keys()
if item.endswith(".png")
],
inputs=[click_image],
label="File Examples",
)
with gr.Tab("Draw"):
sketch = gr.Sketchpad(
sources=(), brush=Brush(colors=["#000000"])
)
with gr.Row():
sketch_submit_btn = gr.Button("Upload", variant="primary")
with gr.Tab("Plan"):
planbot = gr.JSON(elem_classes="json")
with gr.Tab("Memory"):
memory_table = gr.DataFrame(
label="Memory",
headers=["Resource", "Type"],
row_count=5,
wrap=True,
)
chatbot.like(
loop.vote,
[
user_state,
chatbot,
],
[
user_state,
chatbot,
],
)
reply_inputs = [user_state, click_image, chatbot, planbot]
reply_outputs = [user_state, chatbot, memory_table]
add_text = [
partial(loop.add_text, role="human"),
[chatbot, text],
[chatbot, text],
]
text.submit(*add_text).then(loop.plan, reply_inputs, reply_inputs).then(
loop.execute, reply_inputs, reply_inputs
).then(loop.reply, [user_state, chatbot], reply_outputs)
add_msg = [
partial(loop.add_msg, role="human"),
[chatbot, text, audio],
[chatbot, text],
]
submit.click(*add_msg).then(
partial(on_switch_input, disable=True),
[state_input, text, audio],
[state_input, text, audio],
).then(loop.plan, reply_inputs, reply_inputs).then(
loop.execute, reply_inputs, reply_inputs
).then(
loop.reply, [user_state, chatbot], reply_outputs
)
upload_btn.upload(
loop.add_file,
inputs=[user_state, chatbot, upload_btn],
outputs=[user_state, click_image, chatbot, memory_table],
)
record.click(
on_switch_input,
[state_input, text, audio],
[state_input, text, audio],
)
click_image.select(
loop.save_point, [user_state, chatbot], [user_state, chatbot]
).then(loop.plan, reply_inputs, reply_inputs).then(
loop.execute, reply_inputs, reply_inputs
).then(
loop.reply, [user_state, chatbot], reply_outputs
)
click_image.upload(
loop.add_file,
inputs=[user_state, chatbot, click_image],
outputs=[user_state, click_image, chatbot, memory_table],
)
click_image_submit_btn.click(
loop.add_file,
inputs=[user_state, chatbot, click_image],
outputs=[user_state, click_image, chatbot, memory_table],
)
sketch_submit_btn.click(
loop.add_sketch,
inputs=[user_state, chatbot, sketch],
outputs=[user_state, chatbot, memory_table],
)
if https:
demo.queue().launch(
server_name="0.0.0.0",
ssl_certfile="./certificate/cert.pem",
ssl_keyfile="./certificate/key.pem",
ssl_verify=False,
show_api=False,
allowed_paths=[
"assets/human.png",
"assets/assistant.png",
],
**kwargs,
)
else:
demo.queue().launch(
server_name="0.0.0.0",
show_api=False,
allowed_paths=[
"assets/human.png",
"assets/assistant.png",
],
**kwargs,
)
if __name__ == "__main__":
os.makedirs(RESOURCE_ROOT, exist_ok=True)
app(controller="cllm.agents.tog.Controller")