src: add main for imaginer

src/main.py

@@ -0,0 +1,249 @@
+# main.py
+#
+# Copyright 2023 Me
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import sys
+import gi
+
+gi.require_version('Gtk', '4.0')
+gi.require_version('Adw', '1')
+
+from gi.repository import Gtk, Gio, Adw, GLib, Pango
+from .window import ImaginerWindow
+
+import requests
+from os.path import basename, splitext
+import io
+import threading
+import json
+from PIL import Image, UnidentifiedImageError
+import time
+
+from enum import Enum
+import requests
+import openai
+import re
+import unicodedata
+from time import gmtime, strftime
+
+class ImaginerApplication(Adw.Application):
+    """The main application singleton class."""
+
+    def __init__(self):
+        super().__init__(application_id='io.github.ImaginerApp.Imaginer',
+                         flags=Gio.ApplicationFlags.DEFAULT_FLAGS)
+        self.create_action('quit', self.on_quit_action, ['<primary>q'])
+        self.create_action('about', self.on_about_action)
+        self.create_action('preferences', self.on_preferences_action)
+        self.create_action('get_started', self.on_get_started_action)
+        self.create_action('imagine', self.on_imagine_action, ['<primary>Return'])
+        self.create_action('choose_output', self.on_file_chooser, ['<primary>s'])
+        self.create_action('new_window', self.on_new_window_action, ['<primary>n'])
+    
+    def on_quit_action(self, action, _):
+        """Callback for the app.quit action."""
+        self.quit()
+        
+    def do_activate(self):
+        """Called when the application is activated.
+
+        We raise the application's main window, creating it if
+        necessary.
+        """
+        self.win = self.props.active_window
+        if not self.win:
+            self.win = ImaginerWindow(application=self)
+        self.win.present()
+
+        self.file_chooser = Gtk.FileChooserNative()
+        self.file_chooser.set_title(_("Choose a directory"))
+        self.file_chooser.set_transient_for(self.win)
+        self.file_chooser.set_action(Gtk.FileChooserAction.SELECT_FOLDER)
+        self.file_chooser.set_modal(True)
+        self.file_chooser.connect(
+            "response", self.on_file_chooser_response
+        )
+        self.token = ""
+
+    def on_new_window_action(self, action, _):
+        """Callback for the app.new_window action."""
+        ImaginerWindow(application=self).present()
+
+    def query(self, payload, url):
+        if self.token:
+            headers = {"Authorization": f"Bearer {self.token}"}
+        else:
+            headers = {}
+        response = requests.post(url, headers=headers, json=payload)
+        return response.content
+
+    def on_file_chooser(self, widget, _):
+        """Callback for the app.choose_output action."""
+        self.file_chooser.show()
+
+    def on_file_chooser_response(self, _, response):
+        if response == Gtk.ResponseType.ACCEPT:
+            self.directory = self.file_chooser.get_file()
+            dir_basename = self.directory.get_basename()
+            self.win.label_output.set_label(dir_basename)
+            self.win.button_imagine.set_has_tooltip(False)
+
+        self.file_chooser.hide()
+
+        if response == Gtk.ResponseType.ACCEPT:
+            self.file_path = self.directory.get_path()
+
+    def on_about_action(self, widget, _):
+        """Callback for the app.about action."""
+        about = Adw.AboutWindow(transient_for=self.props.active_window,
+                                application_name='imaginer',
+                                application_icon='io.github.ImaginerApp.Imaginer',
+                                developer_name='Me',
+                                version='0.1.0',
+                                developers=['Me'],
+                                copyright='© 2023 Me')
+        about.present()
+
+    def on_get_started_action(self, widget, _):
+        """Callback for the app.get_started action."""
+        self.win.stack_imaginer.set_visible_child_name("stack_imagine")
+
+    def slugify(self, value):
+        value = unicodedata.normalize('NFKD', value).encode('ascii', 'ignore').decode('ascii')
+        value = re.sub('[^\w\s-]', '', value).strip().lower()
+        return re.sub('[-\s]+', '-', value)
+
+    def on_imagine_action(self, widget, _):
+        """Callback for the app.imagine action."""
+        self.win.banner.set_revealed(False)
+        self.win.stack_imaginer.set_visible_child_name("stack_loading")
+        self.win.spinner_loading.start()
+
+        self.provider = self.win.provider.props.selected
+
+        class ProvidersEnum(Enum):
+            STABLE_DIFFUSION = 0
+            OPENAI = 1
+            WAIFU_DIFFUSION = 2
+            OPENJOURNEY = 3
+            NITRO_DIFFUSION = 4
+            ANALOG_DIFFUSION = 5
+            PORTRAIT_PLUS = 6
+        
+        
+
+        prompt = self.win.prompt.get_text()
+        self.token = self.win.token.get_text()
+        openai.api_key = self.token
+
+        try:
+            path = self.file_path
+            print(path)
+            path = f"{path}/imaginer-{self.slugify(prompt)}-{strftime('%d-%b-%Y-%H-%M-%S', gmtime())}.png"
+        except AttributeError:
+            path = "imaginer.png"
+
+        def thread_run():
+            match self.provider:
+                case ProvidersEnum.OPENAI.value:
+                    try:
+                        response = openai.Image.create(
+                            prompt=prompt,
+                            n=1,
+                            size="1024x1024"
+                        )
+                        image_url = response['data'][0]['url']
+                        image_bytes = requests.get(image_url).content
+                    except openai.error.AuthenticationError:
+                        self.win.banner.set_title("Invalid API Key")
+                        self.win.banner.set_revealed(True)
+                        image_bytes = None
+
+                case ProvidersEnum.STABLE_DIFFUSION.value:
+                    image_bytes = self.query({
+                        "inputs": prompt,
+                    }, "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2-1")
+                case ProvidersEnum.WAIFU_DIFFUSION.value:
+                    image_bytes = self.query({
+                        "inputs": prompt,
+                    }, "https://api-inference.huggingface.co/models/hakurei/waifu-diffusion")
+                case ProvidersEnum.OPENJOURNEY.value:
+                    image_bytes = self.query({
+                        "inputs": prompt,
+                    }, "https://api-inference.huggingface.co/models/prompthero/openjourney")
+                case ProvidersEnum.NITRO_DIFFUSION.value:
+                    image_bytes = self.query({
+                        "inputs": prompt,
+                    }, "https://api-inference.huggingface.co/models/nitrosocke/Nitro-Diffusion")
+                case ProvidersEnum.ANALOG_DIFFUSION.value:
+                    image_bytes = self.query({
+                        "inputs": prompt,
+                    }, "https://api-inference.huggingface.co/models/wavymulder/Analog-Diffusion")
+                case ProvidersEnum.PORTRAIT_PLUS.value:
+                    image_bytes = self.query({
+                        "inputs": prompt,
+                    }, "https://api-inference.huggingface.co/models/wavymulder/portraitplus")
+            if image_bytes:
+                try:
+                    image = Image.open(io.BytesIO(image_bytes))
+                except UnidentifiedImageError:
+                    error = json.loads(image_bytes)["error"]                    
+                    self.win.banner.set_title(error)
+                    self.win.banner.set_revealed(True)
+                    image = None
+            else:
+                image = None
+            
+            GLib.idle_add(cleanup, image)
+
+        def cleanup(image):
+            self.win.spinner_loading.stop()
+            self.win.stack_imaginer.set_visible_child_name("stack_imagine")
+            t.join()
+            if image:
+                image.save(path)
+                self.win.image.set_file(Gio.File.new_for_path(path))
+            
+        t = threading.Thread(target=thread_run)
+        t.start()
+
+
+    def on_preferences_action(self, widget, _):
+        """Callback for the app.preferences action."""
+        print('app.preferences action activated')
+
+    def create_action(self, name, callback, shortcuts=None):
+        """Add an application action.
+
+        Args:
+            name: the name of the action
+            callback: the function to be called when the action is
+              activated
+            shortcuts: an optional list of accelerators
+        """
+        action = Gio.SimpleAction.new(name, None)
+        action.connect("activate", callback)
+        self.add_action(action)
+        if shortcuts:
+            self.set_accels_for_action(f"app.{name}", shortcuts)
+
+
+def main(version):
+    """The application's entry point."""
+    app = ImaginerApplication()
+    return app.run(sys.argv)