diff --git a/src/main.py b/src/main.py
new file mode 100644
index 0000000..ecaa9a7
--- /dev/null
+++ b/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 .
+#
+# 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, ['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, ['Return'])
+ self.create_action('choose_output', self.on_file_chooser, ['s'])
+ self.create_action('new_window', self.on_new_window_action, ['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)