Novetat: ara pot generar imatges

mastobot.py

@@ -4,6 +4,9 @@ from mastodon import Mastodon
 import re
 import time
 import sys
+import os
+import requests
+import base64
 import pdb
 
 def cleanhtml(raw_html):
@@ -102,7 +105,44 @@ if __name__ == '__main__':
 
         if reply == True:
 
-            answer = xat.completions(prompt=query_word)
+            prompt = query_word
+
+            if prompt[:7] == 'imatge:':
+
+                try:
+
+                    imatge = xat.create_image(prompt=prompt[7:])
+
+                    if not 'error' in imatge:
+
+                        with open("images/imatge.png", "wb") as fh:
+                            fh.write(base64.b64decode(imatge['data'][0]['b64_json']))
+
+                        post_text = f'@{username} {prompt[8:]}'
+
+                        image_id = mastodon.media_post('images/imatge.png', "image/png", description=prompt[8:]).id
+
+                        mastodon.status_post(post_text, in_reply_to_id=status_id, media_ids={image_id}, visibility=visibility)
+
+                    else:
+
+                        post_text = f'@{username} {prompt[8:]}'
+
+                        post_text += f"\n\n{imatge['error']['message']}"
+
+                        mastodon.status_post(post_text, in_reply_to_id=status_id, visibility=visibility)
+
+                except:
+
+                    pass
+
+                    post_text = f'@{username} {prompt[8:]}\n\n{imatge}'
+
+                    mastodon.status_post(post_text, in_reply_to_id=status_id, visibility=visibility)
+
+            else:
+
+                answer = xat.completions(prompt=prompt)
 
                 answer = answer.choices[0]['text']
 

xatapi.py

@@ -86,6 +86,23 @@ class Xat:
 
         return response
 
+    def create_image(self, model=None, prompt=None):
+
+        data = {
+            'prompt': prompt,
+            'n': 1,
+            'size': "512x512",
+            'response_format': "b64_json"
+            }
+
+        endpoint = f'{self.api_base_url}/v1/images/generations'
+
+        response = self.__api_request('POST', endpoint, data)
+
+        response = self.__json_allow_dict_attrs(response.json())
+
+        return response
+
     @staticmethod
     def __check_setup(self):
 
@@ -165,6 +182,8 @@ class Xat:
                         # this is for compatibility with older versions
                         # which raised OpenaiAPIError('Endpoint not found.')
                         # on any 404
+                elif response.status_code == 400:
+                    return response
                 elif response.status_code == 401:
                     ex_type = OpenaiUnauthorizedError
                 elif response.status_code == 422: