Merge branch 'master' into command_registry

scripts/ai_config.py

@@ -1,6 +1,6 @@
 import yaml
 import data
-
+import os
 
 class AIConfig:
     def __init__(self, ai_name="", ai_role="", ai_goals=[]):
@@ -9,7 +9,7 @@ class AIConfig:
         self.ai_goals = ai_goals
 
     # Soon this will go in a folder where it remembers more stuff about the run(s)
-    SAVE_FILE = "../ai_settings.yaml"
+    SAVE_FILE = os.path.join(os.path.dirname(__file__), '..', 'ai_settings.yaml')
 
     @classmethod
     def load(cls, config_file=SAVE_FILE):

scripts/config.py

@@ -53,6 +53,9 @@ class Config(metaclass=Singleton):
         self.pinecone_api_key = os.getenv("PINECONE_API_KEY")
         self.pinecone_region = os.getenv("PINECONE_ENV")
 
+        self.image_provider = os.getenv("IMAGE_PROVIDER")
+        self.huggingface_api_token = os.getenv("HUGGINGFACE_API_TOKEN")
+
         # User agent headers to use when browsing web
         # Some websites might just completely deny request with an error code if no user agent was found.
         self.user_agent_header = {"User-Agent":"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/83.0.4103.97 Safari/537.36"}

scripts/data/prompt.txt

@@ -23,6 +23,7 @@ COMMANDS:
 17. Write Tests: "write_tests", args: "code": "<full_code_string>", "focus": "<list_of_focus_areas>"
 18. Execute Python File: "execute_python_file", args: "file": "<file>"
 19. Task Complete (Shutdown): "task_complete", args: "reason": "<reason>"
+20. Generate Image: "generate_image", args: "prompt": "<prompt>"
 
 RESOURCES:
 

scripts/image_gen.py

@@ -0,0 +1,58 @@
+import requests
+import io
+import os.path
+from PIL import Image
+from config import Config
+import uuid
+import openai
+from base64 import b64decode
+
+cfg = Config()
+
+working_directory = "auto_gpt_workspace"
+
+@command("generate_image", "Generate Image", '"prompt": "<prompt>"')
+def generate_image(prompt):
+
+    filename = str(uuid.uuid4()) + ".jpg"
+    
+    # DALL-E
+    if cfg.image_provider == 'dalle':
+
+        openai.api_key = cfg.openai_api_key
+
+        response = openai.Image.create(
+            prompt=prompt,
+            n=1,
+            size="256x256",
+            response_format="b64_json",
+        )
+
+        print("Image Generated for prompt:" + prompt)
+
+        image_data = b64decode(response["data"][0]["b64_json"])
+
+        with open(working_directory + "/" + filename, mode="wb") as png:
+            png.write(image_data)
+
+        return "Saved to disk:" + filename
+
+    # STABLE DIFFUSION
+    elif cfg.image_provider == 'sd':
+
+        API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
+        headers = {"Authorization": "Bearer " + cfg.huggingface_api_token}
+
+        response = requests.post(API_URL, headers=headers, json={
+            "inputs": prompt,
+        })
+
+        image = Image.open(io.BytesIO(response.content))
+        print("Image Generated for prompt:" + prompt)
+
+        image.save(os.path.join(working_directory, filename))
+
+        return "Saved to disk:" + filename
+
+    else:
+        return "No Image Provider Set"

scripts/main.py

@@ -280,12 +280,12 @@ next_action_count = 0
 # Make a constant:
 user_input = "Determine which next command to use, and respond using the format specified above:"
 
-
+# raise an exception if pinecone_api_key or region is not provided
+if not cfg.pinecone_api_key or not cfg.pinecone_region: raise Exception("Please provide pinecone_api_key and pinecone_region")
 # Initialize memory and make sure it is empty.
 # this is particularly important for indexing and referencing pinecone memory
 memory = PineconeMemory()
 memory.clear()
-
 print('Using memory of type: ' + memory.__class__.__name__)
 
 # Create a CommandRegistry instance and scan default folder