From 548cca19151ae3e9c3734eaeeec53c5fb059b359 Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Mon, 26 Aug 2024 19:10:24 +0200 Subject: [PATCH 01/14] ci: remove Docker layer caching in GitHub Actions Remove the cache-from and cache-to options from the docker/build-push-action step in the Docker build and push workflow. This change may impact build times but simplifies the workflow configuration. --- .github/workflows/docker-build-push.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/docker-build-push.yml b/.github/workflows/docker-build-push.yml index 050fff6..9a6ba57 100644 --- a/.github/workflows/docker-build-push.yml +++ b/.github/workflows/docker-build-push.yml @@ -48,5 +48,4 @@ jobs: push: true tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} - cache-from: type=gha - cache-to: type=gha,mode=max + From 6d90f16b0141199a4d0db5802bb0ffe3110b14b6 Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Wed, 28 Aug 2024 21:19:34 +0200 Subject: [PATCH 02/14] refactor(gui): improve layout and responsiveness - Restructure main UI layout using grid system - Adjust card sizes and add scrolling for better space utilization - Increase spinner size for better visibility --- src/gui.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/src/gui.py b/src/gui.py index 27bc5f1..7b95427 100644 --- a/src/gui.py +++ b/src/gui.py @@ -87,16 +87,19 @@ def __getattr__(self, name): def setup_ui(self): ui.dark_mode().enable() - with ui.column().classes("w-full max-w-full mx-auto space-y-8"): - with ui.card().classes("w-full"): + with ui.grid(columns=2).classes("w-screen h-full gap-4 px-8"): + with ui.card().classes("col-span-2"): ui.label("Flux LoRA API").classes("text-2xl font-bold mb-4") - with ui.row(): - with ui.column(wrap=False): - self.setup_left_panel() - with ui.column(wrap=False): - self.setup_right_panel() - ui.separator() + + with ui.card().classes("h-[70vh] overflow-auto"): + self.setup_left_panel() + + with ui.card().classes("h-[70vh] overflow-auto"): + self.setup_right_panel() + + with ui.card().classes("col-span-2"): self.setup_bottom_panel() + logger.info("UI setup completed") def setup_left_panel(self): @@ -338,7 +341,7 @@ def update_folder_path(self, e): self.folder_input.value = self.folder_path def setup_right_panel(self): - self.spinner = ui.spinner(type="infinity", size="xl") + self.spinner = ui.spinner(type="infinity", size="150px") self.spinner.visible = False self.gallery_container = ui.column().classes("w-full mt-4") From f4dc38da7391060e96c0e717d0e7f974b7aa453e Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Fri, 30 Aug 2024 00:32:19 +0200 Subject: [PATCH 03/14] refactor(config): migrate to dynaconf for settings - Replace JSON-based settings with dynaconf - Update GUI to use new config system - Add settings.toml for default configuration - Modify user model handling in GUI --- requirements.txt | 3 ++- src/config.py | 7 +++++ src/gui.py | 67 ++++++++++++++++++++--------------------------- src/main.py | 3 +-- src/settings.toml | 18 +++++++++++++ 5 files changed, 56 insertions(+), 42 deletions(-) create mode 100644 src/config.py create mode 100644 src/settings.toml diff --git a/requirements.txt b/requirements.txt index 7d91192..90b2e9d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,5 @@ python-dotenv==1.0.1 token_count==0.2.1 loguru==0.7.2 nicegui==1.4.36 -httpx==0.27.0 +httpx==0.27.2 +dynaconf==3.2.6 diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000..24bee4c --- /dev/null +++ b/src/config.py @@ -0,0 +1,7 @@ +from dynaconf import Dynaconf + +settings = Dynaconf( + envvar_prefix="FLUXLORA", + settings_files=["src/settings.toml", "src/.secrets.toml"], + lowercase_read=False, +) diff --git a/src/gui.py b/src/gui.py index 7b95427..93eef33 100644 --- a/src/gui.py +++ b/src/gui.py @@ -7,6 +7,7 @@ from pathlib import Path import httpx +from config import settings from loguru import logger from nicegui import events, ui @@ -52,10 +53,9 @@ def _open(self, url: str) -> None: class ImageGeneratorGUI: def __init__(self, image_generator): self.image_generator = image_generator - self.settings_file = "settings.json" - self.load_settings() + self.settings = settings self.flux_fine_tune_models = self.image_generator.get_flux_fine_tune_models() - self.user_added_models = self.settings.get("user_added_models", []) + self.user_added_models = self.get_setting("user_models", {}) self._attributes = [ "prompt", @@ -79,6 +79,9 @@ def __init__(self, image_generator): self.setup_ui() logger.info("ImageGeneratorGUI initialized") + def get_setting(self, key, default=None): + return self.settings.get(key, default) + def __getattr__(self, name): if name in self._attributes: return self.__dict__.get(name, None) @@ -103,8 +106,9 @@ def setup_ui(self): logger.info("UI setup completed") def setup_left_panel(self): + print(f"Available settings: {settings.as_dict()}") self.replicate_model_input = ( - ui.input("Replicate Model", value=self.settings.get("replicate_model", "")) + ui.input("Replicate Model", value=settings.get("replicate_model", "")) .classes("w-full") .tooltip("Enter the Replicate model URL or identifier") ) @@ -118,11 +122,13 @@ def setup_left_panel(self): on_change=self.select_flux_model, ) .classes("w-full") - .tooltip("Select Model") + .tooltip("Select Public Model") ) with ui.row().classes("w-full"): - self.new_model_input = ui.input(label="New Model").classes("w-3/4") - ui.button("Add Model", on_click=self.add_user_model).classes("w-1/4") + self.new_model_input = ui.input(label="Add Custom LoRA").classes("w-3/4") + ui.button("Add Custom LoRA Model", on_click=self.add_user_model).classes( + "w-1/4" + ) self.user_models_select = ui.select( options=self.user_added_models, @@ -143,7 +149,7 @@ def setup_left_panel(self): ) self.folder_input = ui.input( - label="Output Folder", value=self.folder_path + label="Output Folder", value=settings.output_folder ).classes("w-full") self.folder_input.on("change", self.update_folder_path) @@ -302,8 +308,8 @@ def setup_left_panel(self): def add_user_model(self): new_model = self.new_model_input.value if new_model and new_model not in self.user_added_models: - self.user_added_models.append(new_model) - self.user_models_select.options = self.user_added_models + self.user_added_models[new_model] = new_model + self.user_models_select.options = list(self.user_added_models.keys()) self.new_model_input.value = "" self.save_settings() ui.notify(f"Model '{new_model}' added successfully", type="positive") @@ -313,8 +319,8 @@ def add_user_model(self): def delete_user_model(self): selected_model = self.user_models_select.value if selected_model in self.user_added_models: - self.user_added_models.remove(selected_model) - self.user_models_select.options = self.user_added_models + del self.user_added_models[selected_model] + self.user_models_select.options = list(self.user_added_models.keys()) self.user_models_select.value = None self.save_settings() ui.notify(f"Model '{selected_model}' deleted successfully", type="positive") @@ -469,35 +475,18 @@ def update_gallery(self, image_paths): ) def load_settings(self): - if os.path.exists(self.settings_file): - with open(self.settings_file, "r") as f: - self.settings = json.load(f) - logger.info("Settings loaded successfully") - else: - self.settings = {} - logger.info("No existing settings found, using defaults") + self.settings = settings.as_dict() + self.user_added_models = list(settings.get("user_models", {}).keys()) + logger.info("No existing settings found, using defaults") def save_settings(self): - settings_to_save = { - "user_added_models": self.user_added_models, - "replicate_model": self.replicate_model_input.value, - "output_folder": self.folder_path, - "flux_model": self.flux_model, - "aspect_ratio": self.aspect_ratio, - "width": self.width, - "height": self.height, - "num_outputs": self.num_outputs, - "lora_scale": self.lora_scale, - "num_inference_steps": self.num_inference_steps, - "guidance_scale": self.guidance_scale, - "seed": self.seed, - "output_format": self.output_format, - "output_quality": self.output_quality, - "disable_safety_checker": self.disable_safety_checker, - "prompt": self.prompt, - } - with open(self.settings_file, "w") as f: - json.dump(settings_to_save, f) + for attr in self._attributes: + self.settings[attr] = getattr(self, attr) + self.settings["user_models"] = self.user_added_models + self.settings["replicate_model"] = self.replicate_model_input.value + self.settings["output_folder"] = self.folder_path + settings.update(self.settings) + settings.write() logger.info("Settings saved successfully") diff --git a/src/main.py b/src/main.py index f92b29f..47927e1 100644 --- a/src/main.py +++ b/src/main.py @@ -1,9 +1,8 @@ import sys +from gui import create_gui from loguru import logger from nicegui import ui - -from gui import create_gui from replicate_api import ImageGenerator logger.remove() diff --git a/src/settings.toml b/src/settings.toml new file mode 100644 index 0000000..6729098 --- /dev/null +++ b/src/settings.toml @@ -0,0 +1,18 @@ +[default] +replicate_model = "" +output_folder = "" +flux_model = "dev" +aspect_ratio = "1:1" +width = 1024 +height = 1024 +num_outputs = 1 +lora_scale = 0.8 +num_inference_steps = 28 +guidance_scale = 3.5 +seed = -1 +output_format = "png" +output_quality = 80 +disable_safety_checker = true +prompt = "" + +[user_models] From 09aaa70b423c8889500722dfb538224d8185543e Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Fri, 30 Aug 2024 00:44:57 +0200 Subject: [PATCH 04/14] feat(ci): add release trigger to Docker workflow - Change trigger from `push` to `release` with `released` event type to only run on published releases --- .github/workflows/docker-build-push.yml | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/.github/workflows/docker-build-push.yml b/.github/workflows/docker-build-push.yml index 1a16c33..128786e 100644 --- a/.github/workflows/docker-build-push.yml +++ b/.github/workflows/docker-build-push.yml @@ -2,12 +2,9 @@ name: Docker Build and Push on: workflow_dispatch: - push: - branches: - - main - paths: - - src/** - - Dockerfile + release: + types: + - released env: REGISTRY: ghcr.io From 4bf5cc53ba538db4f55b14defda1b6d8402efc68 Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Sat, 31 Aug 2024 18:11:31 +0200 Subject: [PATCH 05/14] refactor(settings): improve config management - Move settings.toml to root directory - Implement local settings override with settings.local.toml - Add reset to default functionality in GUI - Update dependencies in requirements.txt - Refactor settings handling in ImageGeneratorGUI class --- requirements.txt | 5 +- src/settings.toml => settings.toml | 10 ++-- src/config.py | 4 +- src/gui.py | 96 ++++++++++++++++++++++-------- 4 files changed, 82 insertions(+), 33 deletions(-) rename src/settings.toml => settings.toml (100%) diff --git a/requirements.txt b/requirements.txt index 90b2e9d..c19949a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,8 @@ -replicate==0.32.0 +replicate==0.32.1 python-dotenv==1.0.1 token_count==0.2.1 loguru==0.7.2 -nicegui==1.4.36 +nicegui==1.4.37 httpx==0.27.2 dynaconf==3.2.6 +toml==0.10.2 diff --git a/src/settings.toml b/settings.toml similarity index 100% rename from src/settings.toml rename to settings.toml index 6729098..10a7400 100644 --- a/src/settings.toml +++ b/settings.toml @@ -1,18 +1,18 @@ [default] -replicate_model = "" -output_folder = "" flux_model = "dev" aspect_ratio = "1:1" -width = 1024 -height = 1024 num_outputs = 1 lora_scale = 0.8 num_inference_steps = 28 guidance_scale = 3.5 -seed = -1 output_format = "png" output_quality = 80 disable_safety_checker = true +width = 1024 +height = 1024 +seed = -1 +output_folder = "" +replicate_model = "" prompt = "" [user_models] diff --git a/src/config.py b/src/config.py index 24bee4c..f4f6add 100644 --- a/src/config.py +++ b/src/config.py @@ -2,6 +2,8 @@ settings = Dynaconf( envvar_prefix="FLUXLORA", - settings_files=["src/settings.toml", "src/.secrets.toml"], + settings_files=["settings.toml", "settings.local.toml", ".secrets.toml"], + environments=True, + load_dotenv=True, lowercase_read=False, ) diff --git a/src/gui.py b/src/gui.py index 93eef33..8d23217 100644 --- a/src/gui.py +++ b/src/gui.py @@ -7,6 +7,7 @@ from pathlib import Path import httpx +import toml from config import settings from loguru import logger from nicegui import events, ui @@ -55,7 +56,7 @@ def __init__(self, image_generator): self.image_generator = image_generator self.settings = settings self.flux_fine_tune_models = self.image_generator.get_flux_fine_tune_models() - self.user_added_models = self.get_setting("user_models", {}) + self.user_added_models = {} self._attributes = [ "prompt", @@ -71,21 +72,23 @@ def __init__(self, image_generator): "width", "height", "seed", + "output_folder", + "replicate_model", ] - for attr in self._attributes: - setattr(self, attr, self.settings.get(attr, None)) + self.load_settings() + + # Set default output folder if not present in settings + if not self.output_folder: + self.output_folder = ( + str(Path.home() / "Downloads") if not DOCKERIZED else "/app/output" + ) self.setup_ui() logger.info("ImageGeneratorGUI initialized") def get_setting(self, key, default=None): - return self.settings.get(key, default) - - def __getattr__(self, name): - if name in self._attributes: - return self.__dict__.get(name, None) - return super().__getattribute__(name) + return getattr(self, key, default) def setup_ui(self): ui.dark_mode().enable() @@ -131,7 +134,7 @@ def setup_left_panel(self): ) self.user_models_select = ui.select( - options=self.user_added_models, + options=list(self.user_added_models.keys()), label="User Added Models", value=None, on_change=self.select_user_model, @@ -149,7 +152,7 @@ def setup_left_panel(self): ) self.folder_input = ui.input( - label="Output Folder", value=settings.output_folder + label="Output Folder", value=self.output_folder ).classes("w-full") self.folder_input.on("change", self.update_folder_path) @@ -336,15 +339,17 @@ def select_user_model(self, e): def update_folder_path(self, e): new_path = e.value if os.path.isdir(new_path): - self.folder_path = new_path + self.output_folder = new_path self.save_settings() - logger.info(f"Output folder set to: {self.folder_path}") - ui.notify(f"Output folder updated to: {self.folder_path}", type="positive") + logger.info(f"Output folder set to: {self.output_folder}") + ui.notify( + f"Output folder updated to: {self.output_folder}", type="positive" + ) else: ui.notify( "Invalid folder path. Please enter a valid directory.", type="negative" ) - self.folder_input.value = self.folder_path + self.folder_input.value = self.output_folder def setup_right_panel(self): self.spinner = ui.spinner(type="infinity", size="150px") @@ -365,11 +370,17 @@ def setup_bottom_panel(self): ).classes( "w-full bg-blue-500 hover:bg-blue-600 text-white font-bold py-2 px-4 rounded" ) + self.reset_button = ui.button( + "Reset to Default", on_click=self.reset_to_default + ).classes( + "w-1/2 bg-gray-500 hover:bg-gray-600 text-white font-bold py-2 px-4 rounded" + ) def update_replicate_model(self, e): new_model = self.replicate_model_input.value if new_model: self.image_generator.set_model(new_model) + self.replicate_model = new_model self.save_settings() logger.info(f"Replicate model updated to: {new_model}") self.generate_button.enable() @@ -393,6 +404,26 @@ def toggle_custom_dimensions(self, e): self.save_settings() logger.info(f"Custom dimensions toggled: {e.value}") + def reset_to_default(self): + with open("settings.toml", "r") as f: + default_settings = toml.load(f)["default"] + + for attr in self._attributes: + if attr in default_settings: + value = default_settings[attr] + setattr(self, attr, value) + if hasattr(self, f"{attr}_input"): + getattr(self, f"{attr}_input").value = value + elif hasattr(self, f"{attr}_select"): + getattr(self, f"{attr}_select").value = value + elif hasattr(self, f"{attr}_switch"): + getattr(self, f"{attr}_switch").value = value + + self.user_added_models = {} + self.save_settings() + ui.notify("Settings reset to default values", type="info") + logger.info("Settings reset to default values") + async def start_generation(self): if not self.replicate_model_input.value: ui.notify( @@ -455,7 +486,7 @@ async def download_and_display_images(self, image_urls): timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") url_part = urllib.parse.urlparse(url).path.split("/")[-2][:8] file_name = f"generated_image_{timestamp}_{url_part}_{i+1}.png" - file_path = Path(self.folder_path) / file_name + file_path = Path(self.output_folder) / file_name with open(file_path, "wb") as f: f.write(response.content) downloaded_images.append(str(file_path)) @@ -475,18 +506,33 @@ def update_gallery(self, image_paths): ) def load_settings(self): - self.settings = settings.as_dict() - self.user_added_models = list(settings.get("user_models", {}).keys()) - logger.info("No existing settings found, using defaults") + # Load default settings + with open("settings.toml", "r") as f: + default_settings = toml.load(f)["default"] + + # Load local settings + local_settings = {} + if os.path.exists("settings.local.toml"): + with open("settings.local.toml", "r") as f: + local_settings = toml.load(f).get("default", {}) + + # Merge settings, prioritizing local settings + for attr in self._attributes: + setattr(self, attr, local_settings.get(attr, default_settings.get(attr))) + + self.user_added_models = local_settings.get("user_models", {}) def save_settings(self): + settings_dict = {} for attr in self._attributes: - self.settings[attr] = getattr(self, attr) - self.settings["user_models"] = self.user_added_models - self.settings["replicate_model"] = self.replicate_model_input.value - self.settings["output_folder"] = self.folder_path - settings.update(self.settings) - settings.write() + settings_dict[attr] = getattr(self, attr) + settings_dict["user_models"] = self.user_added_models + settings_dict["replicate_model"] = self.replicate_model_input.value + + # Save settings to settings.local.toml + with open("settings.local.toml", "w") as f: + toml.dump({"default": settings_dict}, f) + logger.info("Settings saved successfully") From 379930dd1cf93b00a1172ff9fb3c3240bd56b02a Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Sat, 31 Aug 2024 19:15:45 +0200 Subject: [PATCH 06/14] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(gui.py):=20refactor?= =?UTF-8?q?=20to=20load=20flux=20models=20asynchronously=20for=20better=20?= =?UTF-8?q?UI=20responsiveness=20=E2=9C=A8=20(gui.py):=20add=20async=20sup?= =?UTF-8?q?port=20for=20updating=20and=20selecting=20models?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/gui.py | 119 ++++++++++++++++++++++++++++++----------------------- 1 file changed, 67 insertions(+), 52 deletions(-) diff --git a/src/gui.py b/src/gui.py index 8d23217..a2c1ee8 100644 --- a/src/gui.py +++ b/src/gui.py @@ -55,7 +55,7 @@ class ImageGeneratorGUI: def __init__(self, image_generator): self.image_generator = image_generator self.settings = settings - self.flux_fine_tune_models = self.image_generator.get_flux_fine_tune_models() + self.flux_fine_tune_models = [] self.user_added_models = {} self._attributes = [ @@ -87,8 +87,15 @@ def __init__(self, image_generator): self.setup_ui() logger.info("ImageGeneratorGUI initialized") - def get_setting(self, key, default=None): - return getattr(self, key, default) + # Set up a timer to load flux models asynchronously + ui.timer(0.1, self.load_flux_models, once=True) + + async def load_flux_models(self): + self.flux_fine_tune_models = await asyncio.to_thread( + self.image_generator.get_flux_fine_tune_models + ) + self.flux_models_ui.refresh() + logger.info("Flux models loaded asynchronously") def setup_ui(self): ui.dark_mode().enable() @@ -109,24 +116,15 @@ def setup_ui(self): logger.info("UI setup completed") def setup_left_panel(self): - print(f"Available settings: {settings.as_dict()}") self.replicate_model_input = ( - ui.input("Replicate Model", value=settings.get("replicate_model", "")) + ui.input("Replicate Model", value=self.settings.get("replicate_model", "")) .classes("w-full") .tooltip("Enter the Replicate model URL or identifier") ) self.replicate_model_input.on("change", self.update_replicate_model) - self.flux_models_select = ( - ui.select( - options=self.flux_fine_tune_models, - label="Flux Fine-Tune Models", - value=None, - on_change=self.select_flux_model, - ) - .classes("w-full") - .tooltip("Select Public Model") - ) + self.flux_models_ui() + with ui.row().classes("w-full"): self.new_model_input = ui.input(label="Add Custom LoRA").classes("w-3/4") ui.button("Add Custom LoRA Model", on_click=self.add_user_model).classes( @@ -307,40 +305,78 @@ def setup_left_panel(self): .tooltip("Disable safety checker for generated images.") .bind_value(self, "disable_safety_checker") ) + self.reset_button = ui.button( + "Reset to Default", on_click=self.reset_to_default + ).classes( + "w-1/2 bg-gray-500 hover:bg-gray-600 text-white font-bold py-2 px-4 rounded" + ) - def add_user_model(self): + @ui.refreshable + def flux_models_ui(self): + self.flux_models_select = ( + ui.select( + options=self.flux_fine_tune_models, + label="Flux Fine-Tune Models", + value=None, + on_change=self.select_flux_model, + ) + .classes("w-full") + .tooltip("Select Public Model") + ) + if not self.flux_fine_tune_models: + ui.label("Loading Flux models...").classes("text-gray-500") + + async def update_replicate_model(self, e): + new_model = self.replicate_model_input.value + if new_model: + await asyncio.to_thread(self.image_generator.set_model, new_model) + self.replicate_model = new_model + await asyncio.to_thread(self.save_settings) + logger.info(f"Replicate model updated to: {new_model}") + self.generate_button.enable() + else: + logger.warning("Empty Replicate model provided") + self.generate_button.disable() + + async def select_flux_model(self, e): + if e.value: + self.replicate_model_input.value = e.value + await self.update_replicate_model(e) + self.flux_models_select.value = None + + async def add_user_model(self): new_model = self.new_model_input.value if new_model and new_model not in self.user_added_models: self.user_added_models[new_model] = new_model self.user_models_select.options = list(self.user_added_models.keys()) self.new_model_input.value = "" - self.save_settings() + await asyncio.to_thread(self.save_settings) ui.notify(f"Model '{new_model}' added successfully", type="positive") else: ui.notify("Invalid model name or model already exists", type="negative") - def delete_user_model(self): + async def delete_user_model(self): selected_model = self.user_models_select.value if selected_model in self.user_added_models: del self.user_added_models[selected_model] self.user_models_select.options = list(self.user_added_models.keys()) self.user_models_select.value = None - self.save_settings() + await asyncio.to_thread(self.save_settings) ui.notify(f"Model '{selected_model}' deleted successfully", type="positive") else: ui.notify("No model selected for deletion", type="negative") - def select_user_model(self, e): + async def select_user_model(self, e): if e.value: self.replicate_model_input.value = e.value - self.update_replicate_model(e) + await self.update_replicate_model(e) self.user_models_select.value = None - def update_folder_path(self, e): + async def update_folder_path(self, e): new_path = e.value if os.path.isdir(new_path): self.output_folder = new_path - self.save_settings() + await asyncio.to_thread(self.save_settings) logger.info(f"Output folder set to: {self.output_folder}") ui.notify( f"Output folder updated to: {self.output_folder}", type="positive" @@ -370,41 +406,18 @@ def setup_bottom_panel(self): ).classes( "w-full bg-blue-500 hover:bg-blue-600 text-white font-bold py-2 px-4 rounded" ) - self.reset_button = ui.button( - "Reset to Default", on_click=self.reset_to_default - ).classes( - "w-1/2 bg-gray-500 hover:bg-gray-600 text-white font-bold py-2 px-4 rounded" - ) - def update_replicate_model(self, e): - new_model = self.replicate_model_input.value - if new_model: - self.image_generator.set_model(new_model) - self.replicate_model = new_model - self.save_settings() - logger.info(f"Replicate model updated to: {new_model}") - self.generate_button.enable() - else: - logger.warning("Empty Replicate model provided") - self.generate_button.disable() - - def select_flux_model(self, e): - if e.value: - self.replicate_model_input.value = e.value - self.update_replicate_model(e) - self.flux_models_select.value = None - - def toggle_custom_dimensions(self, e): + async def toggle_custom_dimensions(self, e): if e.value == "custom": self.width_input.enable() self.height_input.enable() else: self.width_input.disable() self.height_input.disable() - self.save_settings() + await asyncio.to_thread(self.save_settings) logger.info(f"Custom dimensions toggled: {e.value}") - def reset_to_default(self): + async def reset_to_default(self): with open("settings.toml", "r") as f: default_settings = toml.load(f)["default"] @@ -420,7 +433,7 @@ def reset_to_default(self): getattr(self, f"{attr}_switch").value = value self.user_added_models = {} - self.save_settings() + await asyncio.to_thread(self.save_settings) ui.notify("Settings reset to default values", type="info") logger.info("Settings reset to default values") @@ -435,9 +448,11 @@ async def start_generation(self): ) return - self.image_generator.set_model(self.replicate_model_input.value) + await asyncio.to_thread( + self.image_generator.set_model, self.replicate_model_input.value + ) - self.save_settings() + await asyncio.to_thread(self.save_settings) params = { "prompt": self.prompt, "flux_model": self.flux_model, From 4ccde4b2fa39a628bf30779818eb28fb6ab65d07 Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Sun, 1 Sep 2024 17:47:40 +0200 Subject: [PATCH 07/14] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(gui.py):=20remove?= =?UTF-8?q?=20flux=20fine-tune=20models=20loading=20and=20related=20UI=20e?= =?UTF-8?q?lements=20=E2=9C=A8=20(gui.py):=20add=20user=20model=20manageme?= =?UTF-8?q?nt=20popup=20for=20adding=20and=20deleting=20models=20=E2=99=BB?= =?UTF-8?q?=EF=B8=8F=20(replicate=5Fapi.py):=20remove=20get=5Fflux=5Ffine?= =?UTF-8?q?=5Ftune=5Fmodels=20method=20and=20related=20comments?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/gui.py | 176 ++++++++++++++++++------------------------- src/replicate_api.py | 11 --- 2 files changed, 75 insertions(+), 112 deletions(-) diff --git a/src/gui.py b/src/gui.py index a2c1ee8..94f86d6 100644 --- a/src/gui.py +++ b/src/gui.py @@ -55,7 +55,6 @@ class ImageGeneratorGUI: def __init__(self, image_generator): self.image_generator = image_generator self.settings = settings - self.flux_fine_tune_models = [] self.user_added_models = {} self._attributes = [ @@ -78,7 +77,6 @@ def __init__(self, image_generator): self.load_settings() - # Set default output folder if not present in settings if not self.output_folder: self.output_folder = ( str(Path.home() / "Downloads") if not DOCKERIZED else "/app/output" @@ -87,16 +85,6 @@ def __init__(self, image_generator): self.setup_ui() logger.info("ImageGeneratorGUI initialized") - # Set up a timer to load flux models asynchronously - ui.timer(0.1, self.load_flux_models, once=True) - - async def load_flux_models(self): - self.flux_fine_tune_models = await asyncio.to_thread( - self.image_generator.get_flux_fine_tune_models - ) - self.flux_models_ui.refresh() - logger.info("Flux models loaded asynchronously") - def setup_ui(self): ui.dark_mode().enable() @@ -123,30 +111,19 @@ def setup_left_panel(self): ) self.replicate_model_input.on("change", self.update_replicate_model) - self.flux_models_ui() - - with ui.row().classes("w-full"): - self.new_model_input = ui.input(label="Add Custom LoRA").classes("w-3/4") - ui.button("Add Custom LoRA Model", on_click=self.add_user_model).classes( - "w-1/4" + with ui.row().classes("w-full mt-4"): + self.user_model_select = ( + ui.select( + options=list(self.user_added_models.keys()), + label="User Added Models", + value=None, + on_change=self.select_user_model, + ) + .classes("w-5/6") + .tooltip("Select or manage user-added models") ) - - self.user_models_select = ui.select( - options=list(self.user_added_models.keys()), - label="User Added Models", - value=None, - on_change=self.select_user_model, - ).classes("w-full") - - ui.button("Delete Selected Model", on_click=self.delete_user_model).classes( - "w-full" - ) - - if DOCKERIZED: - self.folder_path = self.settings.get("output_folder", str("/app/output")) - else: - self.folder_path = self.settings.get( - "output_folder", str(Path.home() / "Downloads") + ui.button(icon="add").classes("w-1/6").on( + "click", self.open_user_model_popup ) self.folder_input = ui.input( @@ -162,7 +139,7 @@ def setup_left_panel(self): ) .classes("w-full") .tooltip( - "Which model to run inferences with. the dev model needs around 28 steps but the schnell model only needs around 4 steps." + "Which model to run inferences with. The dev model needs around 28 steps but the schnell model only needs around 4 steps." ) .bind_value(self, "flux_model") ) @@ -311,66 +288,87 @@ def setup_left_panel(self): "w-1/2 bg-gray-500 hover:bg-gray-600 text-white font-bold py-2 px-4 rounded" ) - @ui.refreshable - def flux_models_ui(self): - self.flux_models_select = ( - ui.select( - options=self.flux_fine_tune_models, - label="Flux Fine-Tune Models", - value=None, - on_change=self.select_flux_model, - ) + def setup_right_panel(self): + self.spinner = ui.spinner(type="infinity", size="150px") + self.spinner.visible = False + + self.gallery_container = ui.column().classes("w-full mt-4") + self.lightbox = Lightbox() + + def setup_bottom_panel(self): + self.prompt_input = ( + ui.textarea("Prompt", value=self.settings.get("prompt", "")) .classes("w-full") - .tooltip("Select Public Model") + .bind_value(self, "prompt") + .props("clearable") + ) + self.generate_button = ui.button( + "Generate Images", on_click=self.start_generation + ).classes( + "w-full bg-blue-500 hover:bg-blue-600 text-white font-bold py-2 px-4 rounded" ) - if not self.flux_fine_tune_models: - ui.label("Loading Flux models...").classes("text-gray-500") - async def update_replicate_model(self, e): - new_model = self.replicate_model_input.value - if new_model: - await asyncio.to_thread(self.image_generator.set_model, new_model) - self.replicate_model = new_model - await asyncio.to_thread(self.save_settings) - logger.info(f"Replicate model updated to: {new_model}") - self.generate_button.enable() - else: - logger.warning("Empty Replicate model provided") - self.generate_button.disable() + def open_user_model_popup(self): + with ui.dialog() as dialog, ui.card(): + ui.label("Manage User-Added Models").classes("text-xl font-bold mb-4") + new_model_input = ui.input(label="Add New Model").classes("w-full mb-4") + add_button = ui.button( + "Add Model", + on_click=lambda: self.add_user_model(new_model_input.value, dialog), + ) - async def select_flux_model(self, e): - if e.value: - self.replicate_model_input.value = e.value - await self.update_replicate_model(e) - self.flux_models_select.value = None + ui.label("Current User-Added Models:").classes("mt-4 mb-2") + with ui.column().classes("w-full"): + for model in self.user_added_models: + with ui.row().classes("w-full justify-between items-center"): + ui.label(model) + ui.button( + icon="delete", + on_click=lambda m=model: self.delete_user_model(m, dialog), + ).props("flat round color=red") + + ui.button("Close", on_click=dialog.close).classes("mt-4") + dialog.open() - async def add_user_model(self): - new_model = self.new_model_input.value + async def add_user_model(self, new_model, dialog): if new_model and new_model not in self.user_added_models: self.user_added_models[new_model] = new_model - self.user_models_select.options = list(self.user_added_models.keys()) - self.new_model_input.value = "" + self.user_model_select.options = list(self.user_added_models.keys()) await asyncio.to_thread(self.save_settings) ui.notify(f"Model '{new_model}' added successfully", type="positive") + dialog.close() else: ui.notify("Invalid model name or model already exists", type="negative") - async def delete_user_model(self): - selected_model = self.user_models_select.value - if selected_model in self.user_added_models: - del self.user_added_models[selected_model] - self.user_models_select.options = list(self.user_added_models.keys()) - self.user_models_select.value = None + async def delete_user_model(self, model, dialog): + if model in self.user_added_models: + del self.user_added_models[model] + self.user_model_select.options = list(self.user_added_models.keys()) + if self.user_model_select.value == model: + self.user_model_select.value = None await asyncio.to_thread(self.save_settings) - ui.notify(f"Model '{selected_model}' deleted successfully", type="positive") + ui.notify(f"Model '{model}' deleted successfully", type="positive") + dialog.close() else: - ui.notify("No model selected for deletion", type="negative") + ui.notify("Cannot delete this model", type="negative") async def select_user_model(self, e): if e.value: self.replicate_model_input.value = e.value await self.update_replicate_model(e) - self.user_models_select.value = None + self.user_model_select.value = None + + async def update_replicate_model(self, e): + new_model = self.replicate_model_input.value + if new_model: + await asyncio.to_thread(self.image_generator.set_model, new_model) + self.replicate_model = new_model + await asyncio.to_thread(self.save_settings) + logger.info(f"Replicate model updated to: {new_model}") + self.generate_button.enable() + else: + logger.warning("Empty Replicate model provided") + self.generate_button.disable() async def update_folder_path(self, e): new_path = e.value @@ -387,26 +385,6 @@ async def update_folder_path(self, e): ) self.folder_input.value = self.output_folder - def setup_right_panel(self): - self.spinner = ui.spinner(type="infinity", size="150px") - self.spinner.visible = False - - self.gallery_container = ui.column().classes("w-full mt-4") - self.lightbox = Lightbox() - - def setup_bottom_panel(self): - self.prompt_input = ( - ui.textarea("Prompt", value=self.settings.get("prompt", "")) - .classes("w-full") - .bind_value(self, "prompt") - .props("clearable") - ) - self.generate_button = ui.button( - "Generate Images", on_click=self.start_generation - ).classes( - "w-full bg-blue-500 hover:bg-blue-600 text-white font-bold py-2 px-4 rounded" - ) - async def toggle_custom_dimensions(self, e): if e.value == "custom": self.width_input.enable() @@ -521,17 +499,14 @@ def update_gallery(self, image_paths): ) def load_settings(self): - # Load default settings with open("settings.toml", "r") as f: default_settings = toml.load(f)["default"] - # Load local settings local_settings = {} if os.path.exists("settings.local.toml"): with open("settings.local.toml", "r") as f: local_settings = toml.load(f).get("default", {}) - # Merge settings, prioritizing local settings for attr in self._attributes: setattr(self, attr, local_settings.get(attr, default_settings.get(attr))) @@ -544,7 +519,6 @@ def save_settings(self): settings_dict["user_models"] = self.user_added_models settings_dict["replicate_model"] = self.replicate_model_input.value - # Save settings to settings.local.toml with open("settings.local.toml", "w") as f: toml.dump({"default": settings_dict}, f) diff --git a/src/replicate_api.py b/src/replicate_api.py index bd284e8..9730e9f 100644 --- a/src/replicate_api.py +++ b/src/replicate_api.py @@ -38,10 +38,8 @@ def generate_images(self, params): raise ImageGenerationError(error_message) try: - # Remove the Flux model from params and store it separately flux_model = params.pop("flux_model", "dev") - # Add the Flux model choice to the input parameters params["model"] = flux_model logger.info( @@ -58,15 +56,6 @@ def generate_images(self, params): logger.exception(error_message) raise ImageGenerationError(error_message) - def get_flux_fine_tune_models(self): - try: - collection = replicate.collections.get("flux-fine-tunes") - models = collection.models - return [model.name for model in models] - except Exception as e: - logger.error(f"Error fetching flux-fine-tunes models: {str(e)}") - return [] - class ImageGenerationError(Exception): pass From b3b8803efad6e848f14be1ee1492c936ffd2d4b9 Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Mon, 2 Sep 2024 19:13:14 +0200 Subject: [PATCH 08/14] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(gui.py):=20refactor?= =?UTF-8?q?=20user=20model=20management?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Introduce a single select dropdown for managing Replicate models - Remove separate input field for Replicate model - Integrate user-added models management into the model select dropdown - Update model selection logic to update the selected model - Persist user-added models in the local settings file --- src/gui.py | 85 +++++++++++++++++++++++++++++++----------------------- 1 file changed, 49 insertions(+), 36 deletions(-) diff --git a/src/gui.py b/src/gui.py index 94f86d6..e4d06ff 100644 --- a/src/gui.py +++ b/src/gui.py @@ -75,6 +75,10 @@ def __init__(self, image_generator): "replicate_model", ] + # Initialize attributes + for attr in self._attributes: + setattr(self, attr, None) + self.load_settings() if not self.output_folder: @@ -104,23 +108,16 @@ def setup_ui(self): logger.info("UI setup completed") def setup_left_panel(self): - self.replicate_model_input = ( - ui.input("Replicate Model", value=self.settings.get("replicate_model", "")) - .classes("w-full") - .tooltip("Enter the Replicate model URL or identifier") - ) - self.replicate_model_input.on("change", self.update_replicate_model) - - with ui.row().classes("w-full mt-4"): - self.user_model_select = ( + with ui.row().classes("w-full"): + self.replicate_model_select = ( ui.select( - options=list(self.user_added_models.keys()), - label="User Added Models", - value=None, - on_change=self.select_user_model, + options=self.model_options, + label="Replicate Model", + value=self.replicate_model, + on_change=self.update_replicate_model, ) .classes("w-5/6") - .tooltip("Select or manage user-added models") + .tooltip("Select or manage Replicate models") ) ui.button(icon="add").classes("w-1/6").on( "click", self.open_user_model_popup @@ -310,14 +307,14 @@ def setup_bottom_panel(self): def open_user_model_popup(self): with ui.dialog() as dialog, ui.card(): - ui.label("Manage User-Added Models").classes("text-xl font-bold mb-4") + ui.label("Manage Replicate Models").classes("text-xl font-bold mb-4") new_model_input = ui.input(label="Add New Model").classes("w-full mb-4") add_button = ui.button( "Add Model", on_click=lambda: self.add_user_model(new_model_input.value, dialog), ) - ui.label("Current User-Added Models:").classes("mt-4 mb-2") + ui.label("Current Models:").classes("mt-4 mb-2") with ui.column().classes("w-full"): for model in self.user_added_models: with ui.row().classes("w-full justify-between items-center"): @@ -333,7 +330,16 @@ def open_user_model_popup(self): async def add_user_model(self, new_model, dialog): if new_model and new_model not in self.user_added_models: self.user_added_models[new_model] = new_model - self.user_model_select.options = list(self.user_added_models.keys()) + self.model_options = list(self.user_added_models.keys()) + self.replicate_model_select.options = self.model_options + self.replicate_model_select.value = new_model + await self.update_replicate_model( + events.ValueChangeEventArguments( + sender=self.replicate_model_select, + value=new_model, + client=ui.client, + ) + ) await asyncio.to_thread(self.save_settings) ui.notify(f"Model '{new_model}' added successfully", type="positive") dialog.close() @@ -343,23 +349,23 @@ async def add_user_model(self, new_model, dialog): async def delete_user_model(self, model, dialog): if model in self.user_added_models: del self.user_added_models[model] - self.user_model_select.options = list(self.user_added_models.keys()) - if self.user_model_select.value == model: - self.user_model_select.value = None + self.model_options = list(self.user_added_models.keys()) + self.replicate_model_select.options = self.model_options + if self.replicate_model_select.value == model: + self.replicate_model_select.value = None + await self.update_replicate_model( + events.ValueChangeEventArguments( + sender=self.replicate_model_select, value=None, client=ui.client + ) + ) await asyncio.to_thread(self.save_settings) ui.notify(f"Model '{model}' deleted successfully", type="positive") dialog.close() else: ui.notify("Cannot delete this model", type="negative") - async def select_user_model(self, e): - if e.value: - self.replicate_model_input.value = e.value - await self.update_replicate_model(e) - self.user_model_select.value = None - async def update_replicate_model(self, e): - new_model = self.replicate_model_input.value + new_model = e.value if new_model: await asyncio.to_thread(self.image_generator.set_model, new_model) self.replicate_model = new_model @@ -367,7 +373,7 @@ async def update_replicate_model(self, e): logger.info(f"Replicate model updated to: {new_model}") self.generate_button.enable() else: - logger.warning("Empty Replicate model provided") + logger.warning("No Replicate model selected") self.generate_button.disable() async def update_folder_path(self, e): @@ -416,23 +422,23 @@ async def reset_to_default(self): logger.info("Settings reset to default values") async def start_generation(self): - if not self.replicate_model_input.value: + if not self.replicate_model_select.value: ui.notify( - "Please set a Replicate model before generating images.", + "Please select a Replicate model before generating images.", type="negative", ) logger.warning( - "Attempted to generate images without setting a Replicate model" + "Attempted to generate images without selecting a Replicate model" ) return await asyncio.to_thread( - self.image_generator.set_model, self.replicate_model_input.value + self.image_generator.set_model, self.replicate_model_select.value ) await asyncio.to_thread(self.save_settings) params = { - "prompt": self.prompt, + "prompt": self.prompt_input.value, "flux_model": self.flux_model, "aspect_ratio": self.aspect_ratio, "num_outputs": self.num_outputs, @@ -510,14 +516,21 @@ def load_settings(self): for attr in self._attributes: setattr(self, attr, local_settings.get(attr, default_settings.get(attr))) - self.user_added_models = local_settings.get("user_models", {}) + models = local_settings.get("models", default_settings.get("models", {})) + self.user_added_models = { + model: model for model in models.get("user_added", []) + } + + self.model_options = list(self.user_added_models.keys()) + self.replicate_model = local_settings.get("replicate_model", "") def save_settings(self): settings_dict = {} for attr in self._attributes: settings_dict[attr] = getattr(self, attr) - settings_dict["user_models"] = self.user_added_models - settings_dict["replicate_model"] = self.replicate_model_input.value + + settings_dict["models"] = {"user_added": list(self.user_added_models.keys())} + settings_dict["replicate_model"] = self.replicate_model_select.value with open("settings.local.toml", "w") as f: toml.dump({"default": settings_dict}, f) From ef50a464bea24c0b8c4174a6dae6235775bd2976 Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Mon, 2 Sep 2024 19:13:25 +0200 Subject: [PATCH 09/14] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(settings):=20ignore?= =?UTF-8?q?=20local=20settings=20file?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add `settings.local.toml` to .gitignore to exclude local configuration overrides from version control --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 43b997a..5d62eb2 100644 --- a/.gitignore +++ b/.gitignore @@ -161,3 +161,4 @@ cython_debug/ # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ settings.json +settings.local.toml From b1976d2c2cae3ae28f281f0db5c6530f95073d11 Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Mon, 2 Sep 2024 19:14:09 +0200 Subject: [PATCH 10/14] =?UTF-8?q?=E2=9C=A8=20(settings.toml):=20add=20defa?= =?UTF-8?q?ult=20prompt=20and=20user-added=20models=20configuration?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add a new `models.user_added` configuration to store user-added models --- settings.toml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/settings.toml b/settings.toml index 10a7400..a090177 100644 --- a/settings.toml +++ b/settings.toml @@ -13,6 +13,7 @@ height = 1024 seed = -1 output_folder = "" replicate_model = "" -prompt = "" +prompt = "Enter Prompt here..." -[user_models] +[models] +user_added = [] \ No newline at end of file From 43524ec8eb6fa206e8d4600fb5b1857befd95cc3 Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Mon, 2 Sep 2024 19:21:35 +0200 Subject: [PATCH 11/14] =?UTF-8?q?=E2=9C=A8=20(gui.py):=20add=20confirmatio?= =?UTF-8?q?n=20dialog=20before=20deleting=20user=20model=20=E2=99=BB?= =?UTF-8?q?=EF=B8=8F=20(gui.py):=20refactor=20open=5Fuser=5Fmodel=5Fpopup?= =?UTF-8?q?=20to=20be=20async?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/gui.py | 28 +++++++++++++++++++++++----- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/src/gui.py b/src/gui.py index e4d06ff..588e1c4 100644 --- a/src/gui.py +++ b/src/gui.py @@ -305,11 +305,11 @@ def setup_bottom_panel(self): "w-full bg-blue-500 hover:bg-blue-600 text-white font-bold py-2 px-4 rounded" ) - def open_user_model_popup(self): + async def open_user_model_popup(self): with ui.dialog() as dialog, ui.card(): ui.label("Manage Replicate Models").classes("text-xl font-bold mb-4") new_model_input = ui.input(label="Add New Model").classes("w-full mb-4") - add_button = ui.button( + ui.button( "Add Model", on_click=lambda: self.add_user_model(new_model_input.value, dialog), ) @@ -321,7 +321,9 @@ def open_user_model_popup(self): ui.label(model) ui.button( icon="delete", - on_click=lambda m=model: self.delete_user_model(m, dialog), + on_click=lambda m=model: self.confirm_delete_model( + m, dialog + ), ).props("flat round color=red") ui.button("Close", on_click=dialog.close).classes("mt-4") @@ -346,7 +348,22 @@ async def add_user_model(self, new_model, dialog): else: ui.notify("Invalid model name or model already exists", type="negative") - async def delete_user_model(self, model, dialog): + async def confirm_delete_model(self, model, parent_dialog): + with ui.dialog() as confirm_dialog, ui.card(): + ui.label(f"Are you sure you want to delete the model '{model}'?").classes( + "mb-4" + ) + with ui.row(): + ui.button( + "Yes", + on_click=lambda: self.delete_user_model( + model, parent_dialog, confirm_dialog + ), + ).classes("mr-2") + ui.button("No", on_click=confirm_dialog.close) + confirm_dialog.open() + + async def delete_user_model(self, model, parent_dialog, confirm_dialog): if model in self.user_added_models: del self.user_added_models[model] self.model_options = list(self.user_added_models.keys()) @@ -360,7 +377,8 @@ async def delete_user_model(self, model, dialog): ) await asyncio.to_thread(self.save_settings) ui.notify(f"Model '{model}' deleted successfully", type="positive") - dialog.close() + confirm_dialog.close() + parent_dialog.close() else: ui.notify("Cannot delete this model", type="negative") From ea91b60a62e885ac1c95a36a94add71148438e9a Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Tue, 3 Sep 2024 00:19:50 +0200 Subject: [PATCH 12/14] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(gui.py):=20refactor?= =?UTF-8?q?=20key=20handling=20and=20model=20management=20for=20clarity=20?= =?UTF-8?q?and=20async=20support=20=E2=9C=A8=20(gui.py):=20add=20refreshab?= =?UTF-8?q?le=20model=20list=20and=20async=20save=20settings=20?= =?UTF-8?q?=F0=9F=94=A7=20(gui.py):=20use=20constant=20for=20local=20setti?= =?UTF-8?q?ngs=20file=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/gui.py | 130 +++++++++++++++++++++++++---------------------------- 1 file changed, 62 insertions(+), 68 deletions(-) diff --git a/src/gui.py b/src/gui.py index 588e1c4..112e4a9 100644 --- a/src/gui.py +++ b/src/gui.py @@ -10,7 +10,7 @@ import toml from config import settings from loguru import logger -from nicegui import events, ui +from nicegui import ui logger.remove() logger.add( @@ -20,11 +20,13 @@ DOCKERIZED = os.environ.get("DOCKER_CONTAINER", False) +SETTINGS_LOCAL_FILE = "settings.local.toml" + class Lightbox: def __init__(self): with ui.dialog().props("maximized").classes("bg-black") as self.dialog: - ui.keyboard(self._handle_key) + self.dialog.on_key = self._handle_key self.large_image = ui.image().props("no-spinner fit=scale-down") self.image_list = [] @@ -35,15 +37,15 @@ def add_image(self, thumb_url: str, orig_url: str) -> ui.image: ): return ui.image(thumb_url) - def _handle_key(self, event_args: events.KeyEventArguments) -> None: - if not event_args.action.keydown: + def _handle_key(self, e) -> None: + if not e.action.keydown: return - if event_args.key.escape: + if e.key.escape: self.dialog.close() image_index = self.image_list.index(self.large_image.source) - if event_args.key.arrow_left and image_index > 0: + if e.key.arrow_left and image_index > 0: self._open(self.image_list[image_index - 1]) - if event_args.key.arrow_right and image_index < len(self.image_list) - 1: + if e.key.arrow_right and image_index < len(self.image_list) - 1: self._open(self.image_list[image_index + 1]) def _open(self, url: str) -> None: @@ -86,7 +88,6 @@ def __init__(self, image_generator): str(Path.home() / "Downloads") if not DOCKERIZED else "/app/output" ) - self.setup_ui() logger.info("ImageGeneratorGUI initialized") def setup_ui(self): @@ -114,12 +115,14 @@ def setup_left_panel(self): options=self.model_options, label="Replicate Model", value=self.replicate_model, - on_change=self.update_replicate_model, + on_change=lambda e: asyncio.create_task( + self.update_replicate_model(e.value) + ), ) .classes("w-5/6") .tooltip("Select or manage Replicate models") ) - ui.button(icon="add").classes("w-1/6").on( + ui.button(icon="settings_suggest").classes("w-1/6").on( "click", self.open_user_model_popup ) @@ -280,7 +283,7 @@ def setup_left_panel(self): .bind_value(self, "disable_safety_checker") ) self.reset_button = ui.button( - "Reset to Default", on_click=self.reset_to_default + "Reset Parameters", on_click=self.reset_to_default ).classes( "w-1/2 bg-gray-500 hover:bg-gray-600 text-white font-bold py-2 px-4 rounded" ) @@ -305,89 +308,77 @@ def setup_bottom_panel(self): "w-full bg-blue-500 hover:bg-blue-600 text-white font-bold py-2 px-4 rounded" ) + @ui.refreshable + def model_list(self): + for model in self.user_added_models: + with ui.row().classes("w-full justify-between items-center"): + ui.label(model) + ui.button( + icon="delete", + on_click=lambda m=model: self.confirm_delete_model(m), + ).props("flat round color=red") + async def open_user_model_popup(self): + async def add_model(): + await self.add_user_model(new_model_input.value) + with ui.dialog() as dialog, ui.card(): ui.label("Manage Replicate Models").classes("text-xl font-bold mb-4") new_model_input = ui.input(label="Add New Model").classes("w-full mb-4") - ui.button( - "Add Model", - on_click=lambda: self.add_user_model(new_model_input.value, dialog), - ) + ui.button("Add Model", on_click=add_model) ui.label("Current Models:").classes("mt-4 mb-2") - with ui.column().classes("w-full"): - for model in self.user_added_models: - with ui.row().classes("w-full justify-between items-center"): - ui.label(model) - ui.button( - icon="delete", - on_click=lambda m=model: self.confirm_delete_model( - m, dialog - ), - ).props("flat round color=red") + self.model_list() ui.button("Close", on_click=dialog.close).classes("mt-4") dialog.open() - async def add_user_model(self, new_model, dialog): + async def add_user_model(self, new_model): if new_model and new_model not in self.user_added_models: self.user_added_models[new_model] = new_model self.model_options = list(self.user_added_models.keys()) self.replicate_model_select.options = self.model_options self.replicate_model_select.value = new_model - await self.update_replicate_model( - events.ValueChangeEventArguments( - sender=self.replicate_model_select, - value=new_model, - client=ui.client, - ) - ) - await asyncio.to_thread(self.save_settings) + await self.update_replicate_model(new_model) + await self.save_settings() ui.notify(f"Model '{new_model}' added successfully", type="positive") - dialog.close() + self.model_list.refresh() else: ui.notify("Invalid model name or model already exists", type="negative") - async def confirm_delete_model(self, model, parent_dialog): + async def confirm_delete_model(self, model): + async def delete_model(): + await self.delete_user_model(model, confirm_dialog) + with ui.dialog() as confirm_dialog, ui.card(): ui.label(f"Are you sure you want to delete the model '{model}'?").classes( "mb-4" ) with ui.row(): - ui.button( - "Yes", - on_click=lambda: self.delete_user_model( - model, parent_dialog, confirm_dialog - ), - ).classes("mr-2") + ui.button("Yes", on_click=delete_model).classes("mr-2") ui.button("No", on_click=confirm_dialog.close) confirm_dialog.open() - async def delete_user_model(self, model, parent_dialog, confirm_dialog): + async def delete_user_model(self, model, confirm_dialog): if model in self.user_added_models: del self.user_added_models[model] self.model_options = list(self.user_added_models.keys()) self.replicate_model_select.options = self.model_options if self.replicate_model_select.value == model: self.replicate_model_select.value = None - await self.update_replicate_model( - events.ValueChangeEventArguments( - sender=self.replicate_model_select, value=None, client=ui.client - ) - ) - await asyncio.to_thread(self.save_settings) + await self.update_replicate_model(None) + await self.save_settings() ui.notify(f"Model '{model}' deleted successfully", type="positive") confirm_dialog.close() - parent_dialog.close() + self.model_list.refresh() else: ui.notify("Cannot delete this model", type="negative") - async def update_replicate_model(self, e): - new_model = e.value + async def update_replicate_model(self, new_model): if new_model: await asyncio.to_thread(self.image_generator.set_model, new_model) self.replicate_model = new_model - await asyncio.to_thread(self.save_settings) + await self.save_settings() logger.info(f"Replicate model updated to: {new_model}") self.generate_button.enable() else: @@ -398,7 +389,7 @@ async def update_folder_path(self, e): new_path = e.value if os.path.isdir(new_path): self.output_folder = new_path - await asyncio.to_thread(self.save_settings) + await self.save_settings() logger.info(f"Output folder set to: {self.output_folder}") ui.notify( f"Output folder updated to: {self.output_folder}", type="positive" @@ -416,7 +407,7 @@ async def toggle_custom_dimensions(self, e): else: self.width_input.disable() self.height_input.disable() - await asyncio.to_thread(self.save_settings) + await self.save_settings() logger.info(f"Custom dimensions toggled: {e.value}") async def reset_to_default(self): @@ -424,7 +415,7 @@ async def reset_to_default(self): default_settings = toml.load(f)["default"] for attr in self._attributes: - if attr in default_settings: + if attr in default_settings and attr not in ["models", "replicate_model"]: value = default_settings[attr] setattr(self, attr, value) if hasattr(self, f"{attr}_input"): @@ -434,10 +425,9 @@ async def reset_to_default(self): elif hasattr(self, f"{attr}_switch"): getattr(self, f"{attr}_switch").value = value - self.user_added_models = {} - await asyncio.to_thread(self.save_settings) - ui.notify("Settings reset to default values", type="info") - logger.info("Settings reset to default values") + await self.save_settings() + ui.notify("Parameters reset to default values", type="info") + logger.info("Parameters reset to default values") async def start_generation(self): if not self.replicate_model_select.value: @@ -450,11 +440,12 @@ async def start_generation(self): ) return + # Run set_model in a separate thread if it's a blocking operation await asyncio.to_thread( self.image_generator.set_model, self.replicate_model_select.value ) - await asyncio.to_thread(self.save_settings) + await self.save_settings() params = { "prompt": self.prompt_input.value, "flux_model": self.flux_model, @@ -481,6 +472,7 @@ async def start_generation(self): logger.info(f"Generating images with params: {json.dumps(params, indent=2)}") try: + # Assuming generate_images is also a blocking operation, we'll run it in a thread too output = await asyncio.to_thread( self.image_generator.generate_images, params ) @@ -511,10 +503,10 @@ async def download_and_display_images(self, image_urls): else: logger.error(f"Failed to download image from {url}") - self.update_gallery(downloaded_images) + await self.update_gallery(downloaded_images) ui.notify("Images generated and downloaded successfully!", type="positive") - def update_gallery(self, image_paths): + async def update_gallery(self, image_paths): self.gallery_container.clear() with self.gallery_container: for image_path in image_paths: @@ -527,8 +519,8 @@ def load_settings(self): default_settings = toml.load(f)["default"] local_settings = {} - if os.path.exists("settings.local.toml"): - with open("settings.local.toml", "r") as f: + if os.path.exists(SETTINGS_LOCAL_FILE): + with open(SETTINGS_LOCAL_FILE, "r") as f: local_settings = toml.load(f).get("default", {}) for attr in self._attributes: @@ -542,7 +534,7 @@ def load_settings(self): self.model_options = list(self.user_added_models.keys()) self.replicate_model = local_settings.get("replicate_model", "") - def save_settings(self): + async def save_settings(self): settings_dict = {} for attr in self._attributes: settings_dict[attr] = getattr(self, attr) @@ -550,11 +542,13 @@ def save_settings(self): settings_dict["models"] = {"user_added": list(self.user_added_models.keys())} settings_dict["replicate_model"] = self.replicate_model_select.value - with open("settings.local.toml", "w") as f: + with open(SETTINGS_LOCAL_FILE, "w") as f: toml.dump({"default": settings_dict}, f) logger.info("Settings saved successfully") async def create_gui(image_generator): - return ImageGeneratorGUI(image_generator) + gui = ImageGeneratorGUI(image_generator) + gui.setup_ui() + return gui From 9ccc13952a0f384d937f2975503b916ef70441bd Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Tue, 3 Sep 2024 02:24:47 +0200 Subject: [PATCH 13/14] feat(api-key): add API key management and validation Implement API key input in settings, validation before image generation, and secure storage in .secrets.toml. Update GUI to include API key setup prompts and error handling. --- .gitignore | 1 + src/config.py | 4 +++ src/gui.py | 70 +++++++++++++++++++++++++++++++++++++++----- src/main.py | 9 +++++- src/replicate_api.py | 17 ++++++++++- 5 files changed, 91 insertions(+), 10 deletions(-) diff --git a/.gitignore b/.gitignore index 5d62eb2..8d1c113 100644 --- a/.gitignore +++ b/.gitignore @@ -162,3 +162,4 @@ cython_debug/ #.idea/ settings.json settings.local.toml +.secrets.toml diff --git a/src/config.py b/src/config.py index f4f6add..3eb40a0 100644 --- a/src/config.py +++ b/src/config.py @@ -7,3 +7,7 @@ load_dotenv=True, lowercase_read=False, ) + + +def get_api_key(): + return settings.get("REPLICATE_API_KEY") or settings.get("replicate_api_key") diff --git a/src/gui.py b/src/gui.py index 112e4a9..bd5a82b 100644 --- a/src/gui.py +++ b/src/gui.py @@ -8,7 +8,8 @@ import httpx import toml -from config import settings +from config import get_api_key, settings +from dynaconf import loaders from loguru import logger from nicegui import ui @@ -58,7 +59,7 @@ def __init__(self, image_generator): self.image_generator = image_generator self.settings = settings self.user_added_models = {} - + self.api_key = get_api_key() or os.environ.get("REPLICATE_API_KEY", "") self._attributes = [ "prompt", "flux_model", @@ -92,11 +93,11 @@ def __init__(self, image_generator): def setup_ui(self): ui.dark_mode().enable() + self.check_api_key() with ui.grid(columns=2).classes("w-screen h-full gap-4 px-8"): - with ui.card().classes("col-span-2"): - ui.label("Flux LoRA API").classes("text-2xl font-bold mb-4") - + with ui.card().classes("col-span-full"): + self.setup_top_panel() with ui.card().classes("h-[70vh] overflow-auto"): self.setup_left_panel() @@ -108,8 +109,15 @@ def setup_ui(self): logger.info("UI setup completed") + def setup_top_panel(self): + with ui.card().classes("w-full"): + ui.label("Flux LoRA API").classes("text-2xl font-bold") + ui.button( + icon="settings_suggest", on_click=self.open_settings_popup + ).classes("absolute-right") + def setup_left_panel(self): - with ui.row().classes("w-full"): + with ui.row().classes("w-full items-end"): self.replicate_model_select = ( ui.select( options=self.model_options, @@ -119,10 +127,10 @@ def setup_left_panel(self): self.update_replicate_model(e.value) ), ) - .classes("w-5/6") + .classes("w-4/6 overflow-hidden") .tooltip("Select or manage Replicate models") ) - ui.button(icon="settings_suggest").classes("w-1/6").on( + ui.button(icon="settings_suggest").classes("ml-2").on( "click", self.open_user_model_popup ) @@ -308,6 +316,38 @@ def setup_bottom_panel(self): "w-full bg-blue-500 hover:bg-blue-600 text-white font-bold py-2 px-4 rounded" ) + async def open_settings_popup(self): + with ui.dialog() as dialog, ui.card(): + ui.label("Settings").classes("text-2xl font-bold") + api_key_input = ui.input( + label="API Key", + placeholder="Enter Replicate API Key...", + password_toggle_button=True, + value=self.api_key, + ).classes("w-full") + + async def save_settings(): + new_api_key = api_key_input.value + if new_api_key != self.api_key: + self.api_key = new_api_key + await self.save_api_key() + dialog.close() + ui.notify("Settings saved successfully", type="positive") + + ui.button("Save Settings", on_click=save_settings).classes("mt-4") + dialog.open() + + async def save_api_key(self): + settings.set("REPLICATE_API_KEY", self.api_key) + + secrets_dict = {"default": {"REPLICATE_API_KEY": self.api_key}} + + loaders.write(".secrets.toml", secrets_dict) + + os.environ["REPLICATE_API_KEY"] = self.api_key + + self.image_generator.set_api_key(self.api_key) + @ui.refreshable def model_list(self): for model in self.user_added_models: @@ -410,6 +450,15 @@ async def toggle_custom_dimensions(self, e): await self.save_settings() logger.info(f"Custom dimensions toggled: {e.value}") + def check_api_key(self): + if not self.api_key: + ui.notify( + "No Replicate API Key found. Please set it in the settings before generating images.", + type="warning", + close_button="OK", + timeout=10000, # 10 seconds + ) + async def reset_to_default(self): with open("settings.toml", "r") as f: default_settings = toml.load(f)["default"] @@ -430,6 +479,11 @@ async def reset_to_default(self): logger.info("Parameters reset to default values") async def start_generation(self): + if not self.api_key: + ui.notify( + "Please set your Replicate API Key in the settings.", type="negative" + ) + return if not self.replicate_model_select.value: ui.notify( "Please select a Replicate model before generating images.", diff --git a/src/main.py b/src/main.py index 47927e1..3abd8d0 100644 --- a/src/main.py +++ b/src/main.py @@ -1,5 +1,6 @@ import sys +from config import get_api_key from gui import create_gui from loguru import logger from nicegui import ui @@ -18,6 +19,12 @@ generator = ImageGenerator() +api_key = get_api_key() +if api_key: + generator.set_api_key(api_key) +else: + logger.warning("No Replicate API Key found. Please set it in the settings.") + logger.info("Creating and setting up GUI") @@ -29,4 +36,4 @@ async def main_page(): logger.info("Starting NiceGUI server") -ui.run(title="Replicate Flux LoRA", port=8080) +ui.run(title="Replicate Flux LoRA", port=8080, favicon="🚀") diff --git a/src/replicate_api.py b/src/replicate_api.py index 9730e9f..98bd09d 100644 --- a/src/replicate_api.py +++ b/src/replicate_api.py @@ -1,4 +1,5 @@ import json +import os import sys import replicate @@ -23,8 +24,14 @@ class ImageGenerator: def __init__(self): self.replicate_model = None + self.api_key = None logger.info("ImageGenerator initialized") + def set_api_key(self, api_key): + self.api_key = api_key + os.environ["REPLICATE_API_KEY"] = api_key + logger.info("API key set") + def set_model(self, replicate_model): self.replicate_model = replicate_model logger.info(f"Model set to: {replicate_model}") @@ -37,6 +44,13 @@ def generate_images(self, params): logger.error(error_message) raise ImageGenerationError(error_message) + if not self.api_key: + error_message = ( + "No API key set. Please set an API key before generating images." + ) + logger.error(error_message) + raise ImageGenerationError(error_message) + try: flux_model = params.pop("flux_model", "dev") @@ -47,7 +61,8 @@ def generate_images(self, params): ) logger.info(f"Using Replicate model: {self.replicate_model}") - output = replicate.run(self.replicate_model, input=params) + client = replicate.Client(api_token=self.api_key) + output = client.run(self.replicate_model, input=params) logger.success(f"Images generated successfully. Output: {output}") return output From 1b2b013c0f6e644363bfb83b0c9bee1f35ca409a Mon Sep 17 00:00:00 2001 From: Robin Tuszik <47579899+rtuszik@users.noreply.github.com> Date: Tue, 3 Sep 2024 02:27:13 +0200 Subject: [PATCH 14/14] refactor(gui): optimize ImageGeneratorGUI class Remove unnecessary comments and adjust notification settings --- src/gui.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/gui.py b/src/gui.py index bd5a82b..b5e41c9 100644 --- a/src/gui.py +++ b/src/gui.py @@ -78,7 +78,6 @@ def __init__(self, image_generator): "replicate_model", ] - # Initialize attributes for attr in self._attributes: setattr(self, attr, None) @@ -457,6 +456,7 @@ def check_api_key(self): type="warning", close_button="OK", timeout=10000, # 10 seconds + position="top", ) async def reset_to_default(self): @@ -494,7 +494,6 @@ async def start_generation(self): ) return - # Run set_model in a separate thread if it's a blocking operation await asyncio.to_thread( self.image_generator.set_model, self.replicate_model_select.value ) @@ -526,7 +525,6 @@ async def start_generation(self): logger.info(f"Generating images with params: {json.dumps(params, indent=2)}") try: - # Assuming generate_images is also a blocking operation, we'll run it in a thread too output = await asyncio.to_thread( self.image_generator.generate_images, params )