new

app.py

@@ -1,30 +1,21 @@
-from gradio_tool import GradioTool
 import os
+from langchain.agents import initialize_agent
+from langchain.llms import OpenAI
+from gradio_tools import (StableDiffusionTool, ImageCaptioningTool, StableDiffusionPromptGeneratorTool,
+                          TextToVideoTool)
+from langchain.memory import ConversationBufferMemory
 
-class StableDiffusionTool(GradioTool):
-    """Tool for calling stable diffusion from llm"""
+# Set the API key
+os.environ["OPENAI_API_KEY"] = "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"
 
-    def __init__(
-        self,
-        name="StableDiffusion",
-        description=(
-            "An image generator. Use this to generate images based on "
-            "text input. Input should be a description of what the image should "
-            "look like. The output will be a path to an image file."
-        ),
-        src="gradio-client-demos/stable-diffusion",
-        hf_token=None,
-    ) -> None:
-        super().__init__(name, description, src, hf_token)
+llm = OpenAI(api_key=os.getenv("OPENAI_API_KEY"), temperature=0)
+memory = ConversationBufferMemory(memory_key="chat_history")
+tools = [StableDiffusionTool().langchain, ImageCaptioningTool().langchain,
+         StableDiffusionPromptGeneratorTool().langchain, TextToVideoTool().langchain]
 
-    def create_job(self, query: str) -> Job:
-        return self.client.submit(query, "", 9, fn_index=1)
 
-    def postprocess(self, output: str) -> str:
-        return [os.path.join(output, i) for i in os.listdir(output) if not i.endswith("json")][0]
-
-    def _block_input(self, gr) -> "gr.components.Component":
-        return gr.Textbox()
-
-    def _block_output(self, gr) -> "gr.components.Component":
-        return gr.Image()
+agent = initialize_agent(tools, llm, memory=memory, agent="conversational-react-description", verbose=True)
+output = agent.run(input=("Please create a photo of a dog riding a skateboard "
+                          "but improve my prompt prior to using an image generator."
+                          "Please caption the generated image and create a video for it using the improved prompt."))
+print(output)