>")
+ self.next_btn = QToolButton()
+ self.next_btn.setText(">")
+ self.next_btn.clicked.connect(self.on_show_next_image)
+ self.prompt = QTextEdit()
+ self.prompt.setPlaceholderText("A fantasy landscape")
+ self.prompt.setAcceptRichText(False)
+ self.prompt.setFixedHeight(40)
+ self.neg_prompt = QTextEdit()
+ self.neg_prompt.setPlaceholderText("")
+ self.neg_prompt.setAcceptRichText(False)
+ self.neg_prompt_label = QLabel("Negative prompt (Set guidance scale > 1.0):")
+ self.neg_prompt.setFixedHeight(35)
+ self.neg_prompt.setEnabled(False)
+ self.generate = QPushButton("Generate")
+ self.generate.clicked.connect(self.generate_click)
+ self.browse_results = QPushButton("...")
+ self.browse_results.setFixedWidth(30)
+ self.browse_results.clicked.connect(self.on_open_results_folder)
+ self.browse_results.setToolTip("Open output folder")
+
+ # Create the image navigation layout
+ ilayout = QHBoxLayout()
+ ilayout.addWidget(self.prev_btn)
+ ilayout.addWidget(self.img)
+ ilayout.addWidget(self.next_btn)
+
+ # Create the generate button layout
+ hlayout = QHBoxLayout()
+ hlayout.addWidget(self.neg_prompt)
+ hlayout.addWidget(self.generate)
+ hlayout.addWidget(self.browse_results)
+
+ # Create the actual widget layout
+ vlayout = QVBoxLayout()
+ vlayout.addLayout(ilayout)
+ # vlayout.addItem(self.vspacer)
+ vlayout.addWidget(self.prompt)
+ vlayout.addWidget(self.neg_prompt_label)
+ vlayout.addLayout(hlayout)
+ self.setLayout(vlayout)
+
+ self.parent.settings_changed.connect(self.on_settings_changed)
+
+ def generate_image(self):
+ self.parent.prepare_generation_settings(self.config)
+ self.config.settings.lcm_diffusion_setting.prompt = self.prompt.toPlainText()
+ self.config.settings.lcm_diffusion_setting.negative_prompt = (
+ self.neg_prompt.toPlainText()
+ )
+ images = self.parent.context.generate_text_to_image(
+ self.config.settings,
+ self.config.reshape_required,
+ DEVICE,
+ )
+ self.parent.context.save_images(
+ images,
+ self.config.settings,
+ )
+ self.prepare_images(images)
+ self.after_generation()
+
+ def prepare_images(self, images):
+ """Prepares the generated images to be displayed in the Qt widget"""
+ self.image_index = 0
+ self.gen_images = []
+ for img in images:
+ im = ImageQt(img).copy()
+ pixmap = QPixmap.fromImage(im)
+ self.gen_images.append(pixmap)
+
+ if len(self.gen_images) > 1:
+ self.next_btn.setEnabled(True)
+ self.prev_btn.setEnabled(False)
+ else:
+ self.next_btn.setEnabled(False)
+ self.prev_btn.setEnabled(False)
+
+ self.img.show_image(pixmap=self.gen_images[0])
+
+ def on_show_next_image(self):
+ if self.image_index != len(self.gen_images) - 1 and len(self.gen_images) > 0:
+ self.prev_btn.setEnabled(True)
+ self.image_index += 1
+ self.img.show_image(pixmap=self.gen_images[self.image_index])
+ if self.image_index == len(self.gen_images) - 1:
+ self.next_btn.setEnabled(False)
+
+ def on_show_previous_image(self):
+ if self.image_index != 0:
+ self.next_btn.setEnabled(True)
+ self.image_index -= 1
+ self.img.show_image(pixmap=self.gen_images[self.image_index])
+ if self.image_index == 0:
+ self.prev_btn.setEnabled(False)
+
+ def on_open_results_folder(self):
+ QDesktopServices.openUrl(
+ QUrl.fromLocalFile(self.config.settings.generated_images.path)
+ )
+
+ def generate_click(self):
+ self.img.setText("Please wait...")
+ self.before_generation()
+ worker = ImageGeneratorWorker(self.generate_image)
+ self.parent.threadpool.start(worker)
+
+ def before_generation(self):
+ """Call this function before running a generation task"""
+ self.img.setEnabled(False)
+ self.generate.setEnabled(False)
+ self.browse_results.setEnabled(False)
+
+ def after_generation(self):
+ """Call this function after running a generation task"""
+ self.img.setEnabled(True)
+ self.generate.setEnabled(True)
+ self.browse_results.setEnabled(True)
+ self.parent.store_dimension_settings()
+
+ def on_settings_changed(self):
+ self.neg_prompt.setEnabled(
+ self.config.settings.lcm_diffusion_setting.use_openvino
+ or self.config.settings.lcm_diffusion_setting.use_lcm_lora
+ )
diff --git a/frontend/gui/image_generator_worker.py b/frontend/gui/image_generator_worker.py
new file mode 100644
index 0000000000000000000000000000000000000000..3a948365085ece82337309ac91d278e77fa03e40
--- /dev/null
+++ b/frontend/gui/image_generator_worker.py
@@ -0,0 +1,37 @@
+from PyQt5.QtCore import (
+ pyqtSlot,
+ QRunnable,
+ pyqtSignal,
+ pyqtSlot,
+)
+from PyQt5.QtCore import QObject
+import traceback
+import sys
+
+
+class WorkerSignals(QObject):
+ finished = pyqtSignal()
+ error = pyqtSignal(tuple)
+ result = pyqtSignal(object)
+
+
+class ImageGeneratorWorker(QRunnable):
+ def __init__(self, fn, *args, **kwargs):
+ super(ImageGeneratorWorker, self).__init__()
+ self.fn = fn
+ self.args = args
+ self.kwargs = kwargs
+ self.signals = WorkerSignals()
+
+ @pyqtSlot()
+ def run(self):
+ try:
+ result = self.fn(*self.args, **self.kwargs)
+ except:
+ traceback.print_exc()
+ exctype, value = sys.exc_info()[:2]
+ self.signals.error.emit((exctype, value, traceback.format_exc()))
+ else:
+ self.signals.result.emit(result)
+ finally:
+ self.signals.finished.emit()
diff --git a/frontend/gui/image_variations_widget.py b/frontend/gui/image_variations_widget.py
new file mode 100644
index 0000000000000000000000000000000000000000..b71a7f6a6ba6319664edca595ef279195746c328
--- /dev/null
+++ b/frontend/gui/image_variations_widget.py
@@ -0,0 +1,35 @@
+from PIL import Image
+from PyQt5.QtWidgets import QApplication
+
+from app_settings import AppSettings
+from backend.models.lcmdiffusion_setting import DiffusionTask
+from frontend.gui.img2img_widget import Img2ImgWidget
+from frontend.webui.image_variations_ui import generate_image_variations
+
+
+class ImageVariationsWidget(Img2ImgWidget):
+ def __init__(self, config: AppSettings, parent):
+ super().__init__(config, parent)
+ # Hide prompt and negative prompt widgets
+ self.prompt.hide()
+ self.neg_prompt_label.hide()
+ self.neg_prompt.setEnabled(False)
+
+ def generate_image(self):
+ self.parent.prepare_generation_settings(self.config)
+ self.config.settings.lcm_diffusion_setting.diffusion_task = (
+ DiffusionTask.image_to_image.value
+ )
+ self.config.settings.lcm_diffusion_setting.prompt = ""
+ self.config.settings.lcm_diffusion_setting.negative_prompt = ""
+ self.config.settings.lcm_diffusion_setting.init_image = Image.open(
+ self.img_path.text()
+ )
+ self.config.settings.lcm_diffusion_setting.strength = self.strength.value() / 10
+
+ images = generate_image_variations(
+ self.config.settings.lcm_diffusion_setting.init_image,
+ self.config.settings.lcm_diffusion_setting.strength,
+ )
+ self.prepare_images(images)
+ self.after_generation()
diff --git a/frontend/gui/img2img_widget.py b/frontend/gui/img2img_widget.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e93870270970bfb541f5ec0d414029d1d81dcaa
--- /dev/null
+++ b/frontend/gui/img2img_widget.py
@@ -0,0 +1,115 @@
+from PyQt5.QtWidgets import (
+ QHBoxLayout,
+ QLabel,
+ QLineEdit,
+ QSlider,
+ QToolButton,
+ QFileDialog,
+ QApplication,
+)
+
+
+from PyQt5.QtCore import Qt, QEvent
+
+from PIL import Image
+from constants import DEVICE
+from app_settings import AppSettings
+from urllib.parse import urlparse, unquote
+from frontend.gui.base_widget import BaseWidget
+from backend.models.lcmdiffusion_setting import DiffusionTask
+
+
+class Img2ImgWidget(BaseWidget):
+ def __init__(self, config: AppSettings, parent):
+ super().__init__(config, parent)
+
+ # Create init image selection widgets
+ self.img_label = QLabel("Init image:")
+ self.img_path = QLineEdit()
+ self.img_path.setReadOnly(True)
+ self.img_path.setAcceptDrops(True)
+ self.img_path.installEventFilter(self)
+ self.img_browse = QToolButton()
+ self.img_browse.setText("...")
+ self.img_browse.setToolTip("Browse for an init image")
+ self.img_browse.clicked.connect(self.browse_click)
+ # Create the init image selection layout
+ hlayout = QHBoxLayout()
+ hlayout.addWidget(self.img_label)
+ hlayout.addWidget(self.img_path)
+ hlayout.addWidget(self.img_browse)
+
+ self.strength_label = QLabel("Denoising strength: 0.3")
+ self.strength = QSlider(orientation=Qt.Orientation.Horizontal)
+ self.strength.setMaximum(10)
+ self.strength.setMinimum(1)
+ self.strength.setValue(3)
+ self.strength.valueChanged.connect(self.update_strength_label)
+ # self.layout().insertWidget(1, self.strength_label)
+ # self.layout().insertWidget(2, self.strength)
+ self.layout().addLayout(hlayout)
+ self.layout().addWidget(self.strength_label)
+ self.layout().addWidget(self.strength)
+
+ def browse_click(self):
+ filename = self.show_file_selection_dialog()
+ if filename[0] != "":
+ self.img_path.setText(filename[0])
+
+ def show_file_selection_dialog(self) -> str:
+ filename = QFileDialog.getOpenFileName(
+ self, "Open Image", "results", "Image Files (*.png *.jpg *.bmp)"
+ )
+ return filename
+
+ def eventFilter(self, source, event: QEvent):
+ """This is the Drag and Drop event filter for the init image QLineEdit"""
+ if event.type() == QEvent.DragEnter:
+ if event.mimeData().hasFormat("text/plain"):
+ event.acceptProposedAction()
+ return True
+ elif event.type() == QEvent.Drop:
+ event.acceptProposedAction()
+ path = unquote(urlparse(event.mimeData().text()).path)
+ self.img_path.setText(path)
+ return True
+
+ return False
+
+ def before_generation(self):
+ super().before_generation()
+ self.img_browse.setEnabled(False)
+ self.img_path.setEnabled(False)
+
+ def after_generation(self):
+ super().after_generation()
+ self.img_browse.setEnabled(True)
+ self.img_path.setEnabled(True)
+
+ def generate_image(self):
+ self.parent.prepare_generation_settings(self.config)
+ self.config.settings.lcm_diffusion_setting.diffusion_task = (
+ DiffusionTask.image_to_image.value
+ )
+ self.config.settings.lcm_diffusion_setting.prompt = self.prompt.toPlainText()
+ self.config.settings.lcm_diffusion_setting.negative_prompt = (
+ self.neg_prompt.toPlainText()
+ )
+ self.config.settings.lcm_diffusion_setting.init_image = Image.open(
+ self.img_path.text()
+ )
+ self.config.settings.lcm_diffusion_setting.strength = self.strength.value() / 10
+
+ images = self.parent.context.generate_text_to_image(
+ self.config.settings,
+ self.config.reshape_required,
+ DEVICE,
+ )
+ self.parent.context.save_images(images, self.config.settings)
+ self.prepare_images(images)
+ self.after_generation()
+
+ def update_strength_label(self, value):
+ val = round(int(value) / 10, 1)
+ self.strength_label.setText(f"Denoising strength: {val}")
+ self.config.settings.lcm_diffusion_setting.strength = val
diff --git a/frontend/gui/ui.py b/frontend/gui/ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..9250bf676da1f3dc8a2f5435095b9cec9b08041e
--- /dev/null
+++ b/frontend/gui/ui.py
@@ -0,0 +1,15 @@
+from typing import List
+from frontend.gui.app_window import MainWindow
+from PyQt5.QtWidgets import QApplication
+import sys
+from app_settings import AppSettings
+
+
+def start_gui(
+ argv: List[str],
+ app_settings: AppSettings,
+):
+ app = QApplication(sys.argv)
+ window = MainWindow(app_settings)
+ window.show()
+ app.exec()
diff --git a/frontend/gui/upscaler_widget.py b/frontend/gui/upscaler_widget.py
new file mode 100644
index 0000000000000000000000000000000000000000..c1c8341fb9f185dc63251bd9ced807ed60e76e34
--- /dev/null
+++ b/frontend/gui/upscaler_widget.py
@@ -0,0 +1,92 @@
+from PIL import Image
+from PyQt5.QtWidgets import (
+ QApplication,
+ QHBoxLayout,
+ QRadioButton,
+ QWidget,
+)
+
+from app_settings import AppSettings
+from backend.models.upscale import UpscaleMode
+from backend.upscale.upscaler import upscale_image
+from frontend.gui.img2img_widget import Img2ImgWidget
+from paths import FastStableDiffusionPaths
+
+
+class UpscalerWidget(Img2ImgWidget):
+ scale_factor = 2.0
+
+ def __init__(self, config: AppSettings, parent):
+ super().__init__(config, parent)
+ # Hide prompt and negative prompt widgets
+ self.prompt.hide()
+ self.neg_prompt.hide()
+ # self.neg_prompt.deleteLater()
+ self.strength_label.hide()
+ self.strength.hide()
+ self.generate.setText("Upscale")
+ # Create upscaler widgets
+ self.edsr_button = QRadioButton("EDSR", self)
+ self.edsr_button.toggled.connect(self.on_mode_change)
+ self.edsr_button.toggle()
+ self.sd_button = QRadioButton("SD", self)
+ self.sd_button.toggled.connect(self.on_mode_change)
+ self.aura_button = QRadioButton("AURA-SR", self)
+ self.aura_button.toggled.connect(self.on_mode_change)
+
+ self.neg_prompt_label.setText("Upscale mode (2x) | AURA-SR (4x):")
+ # Create upscaler buttons layout
+ hlayout = QHBoxLayout()
+ hlayout.addWidget(self.edsr_button)
+ hlayout.addWidget(self.sd_button)
+ hlayout.addWidget(self.aura_button)
+ # Can't use a layout with replaceWidget(), so the layout is assigned
+ # to a dummy widget used to replace the negative prompt button and
+ # obtain the desired GUI design
+ self.container = QWidget()
+ self.container.setLayout(hlayout)
+ self.layout().replaceWidget(self.neg_prompt, self.container)
+
+ def generate_image(self):
+ self.parent.prepare_generation_settings(self.config)
+ self.config.settings.lcm_diffusion_setting.init_image = Image.open(
+ self.img_path.text()
+ )
+ self.config.settings.lcm_diffusion_setting.strength = self.strength.value() / 10
+ upscaled_filepath = FastStableDiffusionPaths.get_upscale_filepath(
+ None,
+ self.scale_factor,
+ self.config.settings.generated_images.format,
+ )
+
+ images = upscale_image(
+ context=self.parent.context,
+ src_image_path=self.img_path.text(),
+ dst_image_path=upscaled_filepath,
+ upscale_mode=self.upscale_mode,
+ strength=0.1,
+ )
+ self.prepare_images(images)
+ self.after_generation()
+
+ def before_generation(self):
+ super().before_generation()
+ self.container.setEnabled(False)
+
+ def after_generation(self):
+ super().after_generation()
+ self.container.setEnabled(True)
+ # TODO For some reason, if init_image is not set to None, saving settings
+ # when closing the main window will fail, even though the save() method
+ # explicitly sets init_image to None?
+ self.config.settings.lcm_diffusion_setting.init_image = None
+
+ def on_mode_change(self):
+ self.scale_factor = 2.0
+ if self.edsr_button.isChecked():
+ self.upscale_mode = UpscaleMode.normal.value
+ elif self.sd_button.isChecked():
+ self.upscale_mode = UpscaleMode.sd_upscale.value
+ else:
+ self.upscale_mode = UpscaleMode.aura_sr.value
+ self.scale_factor = 4.0
diff --git a/frontend/utils.py b/frontend/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..61864ce076cf9be32cf383151c4aa3a59540070c
--- /dev/null
+++ b/frontend/utils.py
@@ -0,0 +1,88 @@
+import platform
+from os import path
+from typing import List
+
+from backend.device import is_openvino_device
+from paths import get_file_name
+
+
+def is_reshape_required(
+ prev_width: int,
+ cur_width: int,
+ prev_height: int,
+ cur_height: int,
+ prev_model: int,
+ cur_model: int,
+ prev_num_of_images: int,
+ cur_num_of_images: int,
+) -> bool:
+ reshape_required = False
+ if (
+ prev_width != cur_width
+ or prev_height != cur_height
+ or prev_model != cur_model
+ or prev_num_of_images != cur_num_of_images
+ ):
+ print("Reshape and compile")
+ reshape_required = True
+
+ return reshape_required
+
+
+def enable_openvino_controls() -> bool:
+ return (
+ is_openvino_device()
+ and platform.system().lower() != "darwin"
+ and platform.processor().lower() != "arm"
+ )
+
+
+def get_valid_model_id(
+ models: List,
+ model_id: str,
+ default_model: str = "",
+) -> str:
+ if len(models) == 0:
+ print(
+ "Warning: model configuration file/directory is empty,please add some models."
+ )
+ return ""
+ if model_id == "":
+ if default_model:
+ return default_model
+ else:
+ return models[0]
+
+ if model_id in models:
+ return model_id
+ else:
+ if model_id:
+ print(
+ f"Error:{model_id} Model not found in configuration file,so using first model : {models[0]}"
+ )
+ return models[0]
+
+
+def get_valid_lora_model(
+ models: List,
+ cur_model: str,
+ lora_models_dir: str,
+) -> str:
+ if cur_model == "" or cur_model is None:
+ print(
+ f"No lora models found, please add lora models to {lora_models_dir} directory"
+ )
+ return ""
+ else:
+ if path.exists(cur_model):
+ return get_file_name(cur_model)
+ else:
+ print(f"Lora model {cur_model} not found")
+ if len(models) > 0:
+ print(f"Fallback model - {models[0]}")
+ return get_file_name(models[0])
+ else:
+ print(
+ f"No lora models found, please add lora models to {lora_models_dir} directory"
+ )
+ return ""
diff --git a/frontend/webui/controlnet_ui.py b/frontend/webui/controlnet_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..5bfab4acb318324a97a9f8191ea96feca80e5cd8
--- /dev/null
+++ b/frontend/webui/controlnet_ui.py
@@ -0,0 +1,194 @@
+import gradio as gr
+from PIL import Image
+from backend.lora import get_lora_models
+from state import get_settings
+from backend.models.lcmdiffusion_setting import ControlNetSetting
+from backend.annotators.image_control_factory import ImageControlFactory
+
+_controlnet_models_map = None
+_controlnet_enabled = False
+_adapter_path = None
+
+app_settings = get_settings()
+
+
+def on_user_input(
+ enable: bool,
+ adapter_name: str,
+ conditioning_scale: float,
+ control_image: Image,
+ preprocessor: str,
+):
+ if not isinstance(adapter_name, str):
+ gr.Warning("Please select a valid ControlNet model")
+ return gr.Checkbox(value=False)
+
+ settings = app_settings.settings.lcm_diffusion_setting
+ if settings.controlnet is None:
+ settings.controlnet = ControlNetSetting()
+
+ if enable and (adapter_name is None or adapter_name == ""):
+ gr.Warning("Please select a valid ControlNet adapter")
+ return gr.Checkbox(value=False)
+ elif enable and not control_image:
+ gr.Warning("Please provide a ControlNet control image")
+ return gr.Checkbox(value=False)
+
+ if control_image is None:
+ return gr.Checkbox(value=enable)
+
+ if preprocessor == "None":
+ processed_control_image = control_image
+ else:
+ image_control_factory = ImageControlFactory()
+ control = image_control_factory.create_control(preprocessor)
+ processed_control_image = control.get_control_image(control_image)
+
+ if not enable:
+ settings.controlnet.enabled = False
+ else:
+ settings.controlnet.enabled = True
+ settings.controlnet.adapter_path = _controlnet_models_map[adapter_name]
+ settings.controlnet.conditioning_scale = float(conditioning_scale)
+ settings.controlnet._control_image = processed_control_image
+
+ # This code can be improved; currently, if the user clicks the
+ # "Enable ControlNet" checkbox or changes the currently selected
+ # ControlNet model, it will trigger a pipeline rebuild even if, in
+ # the end, the user leaves the same ControlNet settings
+ global _controlnet_enabled
+ global _adapter_path
+ if settings.controlnet.enabled != _controlnet_enabled or (
+ settings.controlnet.enabled
+ and settings.controlnet.adapter_path != _adapter_path
+ ):
+ settings.rebuild_pipeline = True
+ _controlnet_enabled = settings.controlnet.enabled
+ _adapter_path = settings.controlnet.adapter_path
+ return gr.Checkbox(value=enable)
+
+
+def on_change_conditioning_scale(cond_scale):
+ print(cond_scale)
+ app_settings.settings.lcm_diffusion_setting.controlnet.conditioning_scale = (
+ cond_scale
+ )
+
+
+def get_controlnet_ui() -> None:
+ with gr.Blocks() as ui:
+ gr.HTML(
+ 'Download ControlNet v1.1 model from ControlNet v1.1 (723 MB files) and place it in controlnet_models folder,restart the app'
+ )
+ with gr.Row():
+ with gr.Column():
+ with gr.Row():
+ global _controlnet_models_map
+ _controlnet_models_map = get_lora_models(
+ app_settings.settings.lcm_diffusion_setting.dirs["controlnet"]
+ )
+ controlnet_models = list(_controlnet_models_map.keys())
+ default_model = (
+ controlnet_models[0] if len(controlnet_models) else None
+ )
+
+ enabled_checkbox = gr.Checkbox(
+ label="Enable ControlNet",
+ info="Enable ControlNet",
+ show_label=True,
+ )
+ model_dropdown = gr.Dropdown(
+ _controlnet_models_map.keys(),
+ label="ControlNet model",
+ info="ControlNet model to load (.safetensors format)",
+ value=default_model,
+ interactive=True,
+ )
+ conditioning_scale_slider = gr.Slider(
+ 0.0,
+ 1.0,
+ value=0.5,
+ step=0.05,
+ label="ControlNet conditioning scale",
+ interactive=True,
+ )
+ control_image = gr.Image(
+ label="Control image",
+ type="pil",
+ )
+ preprocessor_radio = gr.Radio(
+ [
+ "Canny",
+ "Depth",
+ "LineArt",
+ "MLSD",
+ "NormalBAE",
+ "Pose",
+ "SoftEdge",
+ "Shuffle",
+ "None",
+ ],
+ label="Preprocessor",
+ info="Select the preprocessor for the control image",
+ value="Canny",
+ interactive=True,
+ )
+
+ enabled_checkbox.input(
+ fn=on_user_input,
+ inputs=[
+ enabled_checkbox,
+ model_dropdown,
+ conditioning_scale_slider,
+ control_image,
+ preprocessor_radio,
+ ],
+ outputs=[enabled_checkbox],
+ )
+ model_dropdown.input(
+ fn=on_user_input,
+ inputs=[
+ enabled_checkbox,
+ model_dropdown,
+ conditioning_scale_slider,
+ control_image,
+ preprocessor_radio,
+ ],
+ outputs=[enabled_checkbox],
+ )
+ conditioning_scale_slider.input(
+ fn=on_user_input,
+ inputs=[
+ enabled_checkbox,
+ model_dropdown,
+ conditioning_scale_slider,
+ control_image,
+ preprocessor_radio,
+ ],
+ outputs=[enabled_checkbox],
+ )
+ control_image.change(
+ fn=on_user_input,
+ inputs=[
+ enabled_checkbox,
+ model_dropdown,
+ conditioning_scale_slider,
+ control_image,
+ preprocessor_radio,
+ ],
+ outputs=[enabled_checkbox],
+ )
+ preprocessor_radio.change(
+ fn=on_user_input,
+ inputs=[
+ enabled_checkbox,
+ model_dropdown,
+ conditioning_scale_slider,
+ control_image,
+ preprocessor_radio,
+ ],
+ outputs=[enabled_checkbox],
+ )
+ conditioning_scale_slider.change(
+ on_change_conditioning_scale, conditioning_scale_slider
+ )
diff --git a/frontend/webui/generation_settings_ui.py b/frontend/webui/generation_settings_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab47e01b3de38f496342aec3931ef8fe202ed709
--- /dev/null
+++ b/frontend/webui/generation_settings_ui.py
@@ -0,0 +1,184 @@
+import gradio as gr
+from state import get_settings
+from backend.models.gen_images import ImageFormat
+
+app_settings = get_settings()
+
+
+def on_change_inference_steps(steps):
+ app_settings.settings.lcm_diffusion_setting.inference_steps = steps
+
+
+def on_change_image_width(img_width):
+ app_settings.settings.lcm_diffusion_setting.image_width = img_width
+
+
+def on_change_image_height(img_height):
+ app_settings.settings.lcm_diffusion_setting.image_height = img_height
+
+
+def on_change_num_images(num_images):
+ app_settings.settings.lcm_diffusion_setting.number_of_images = num_images
+
+
+def on_change_guidance_scale(guidance_scale):
+ app_settings.settings.lcm_diffusion_setting.guidance_scale = guidance_scale
+
+
+def on_change_clip_skip(clip_skip):
+ app_settings.settings.lcm_diffusion_setting.clip_skip = clip_skip
+
+
+def on_change_token_merging(token_merging):
+ app_settings.settings.lcm_diffusion_setting.token_merging = token_merging
+
+
+def on_change_seed_value(seed):
+ app_settings.settings.lcm_diffusion_setting.seed = seed
+
+
+def on_change_seed_checkbox(seed_checkbox):
+ app_settings.settings.lcm_diffusion_setting.use_seed = seed_checkbox
+
+
+def on_change_safety_checker_checkbox(safety_checker_checkbox):
+ app_settings.settings.lcm_diffusion_setting.use_safety_checker = (
+ safety_checker_checkbox
+ )
+
+
+def on_change_tiny_auto_encoder_checkbox(tiny_auto_encoder_checkbox):
+ app_settings.settings.lcm_diffusion_setting.use_tiny_auto_encoder = (
+ tiny_auto_encoder_checkbox
+ )
+
+
+def on_offline_checkbox(offline_checkbox):
+ app_settings.settings.lcm_diffusion_setting.use_offline_model = offline_checkbox
+
+
+def on_change_image_format(image_format):
+ if image_format == "PNG":
+ app_settings.settings.generated_images.format = ImageFormat.PNG.value.upper()
+ else:
+ app_settings.settings.generated_images.format = ImageFormat.JPEG.value.upper()
+
+ app_settings.save()
+
+
+def get_generation_settings_ui() -> None:
+ with gr.Blocks():
+ with gr.Row():
+ with gr.Column():
+ num_inference_steps = gr.Slider(
+ 1,
+ 25,
+ value=app_settings.settings.lcm_diffusion_setting.inference_steps,
+ step=1,
+ label="Inference Steps",
+ interactive=True,
+ )
+
+ image_height = gr.Slider(
+ 256,
+ 1024,
+ value=app_settings.settings.lcm_diffusion_setting.image_height,
+ step=256,
+ label="Image Height",
+ interactive=True,
+ )
+ image_width = gr.Slider(
+ 256,
+ 1024,
+ value=app_settings.settings.lcm_diffusion_setting.image_width,
+ step=256,
+ label="Image Width",
+ interactive=True,
+ )
+ num_images = gr.Slider(
+ 1,
+ 50,
+ value=app_settings.settings.lcm_diffusion_setting.number_of_images,
+ step=1,
+ label="Number of images to generate",
+ interactive=True,
+ )
+ guidance_scale = gr.Slider(
+ 1.0,
+ 10.0,
+ value=app_settings.settings.lcm_diffusion_setting.guidance_scale,
+ step=0.1,
+ label="Guidance Scale",
+ interactive=True,
+ )
+ clip_skip = gr.Slider(
+ 1,
+ 12,
+ value=app_settings.settings.lcm_diffusion_setting.clip_skip,
+ step=1,
+ label="CLIP Skip",
+ interactive=True,
+ )
+ token_merging = gr.Slider(
+ 0.0,
+ 1.0,
+ value=app_settings.settings.lcm_diffusion_setting.token_merging,
+ step=0.01,
+ label="Token Merging",
+ interactive=True,
+ )
+
+
+ seed = gr.Slider(
+ value=app_settings.settings.lcm_diffusion_setting.seed,
+ minimum=0,
+ maximum=999999999,
+ label="Seed",
+ step=1,
+ interactive=True,
+ )
+ seed_checkbox = gr.Checkbox(
+ label="Use seed",
+ value=app_settings.settings.lcm_diffusion_setting.use_seed,
+ interactive=True,
+ )
+
+ safety_checker_checkbox = gr.Checkbox(
+ label="Use Safety Checker",
+ value=app_settings.settings.lcm_diffusion_setting.use_safety_checker,
+ interactive=True,
+ )
+ tiny_auto_encoder_checkbox = gr.Checkbox(
+ label="Use tiny auto encoder for SD",
+ value=app_settings.settings.lcm_diffusion_setting.use_tiny_auto_encoder,
+ interactive=True,
+ )
+ offline_checkbox = gr.Checkbox(
+ label="Use locally cached model or downloaded model folder(offline)",
+ value=app_settings.settings.lcm_diffusion_setting.use_offline_model,
+ interactive=True,
+ )
+ img_format = gr.Radio(
+ label="Output image format",
+ choices=["PNG", "JPEG"],
+ value=app_settings.settings.generated_images.format,
+ interactive=True,
+ )
+
+ num_inference_steps.change(on_change_inference_steps, num_inference_steps)
+ image_height.change(on_change_image_height, image_height)
+ image_width.change(on_change_image_width, image_width)
+ num_images.change(on_change_num_images, num_images)
+ guidance_scale.change(on_change_guidance_scale, guidance_scale)
+ clip_skip.change(on_change_clip_skip, clip_skip)
+ token_merging.change(on_change_token_merging, token_merging)
+ seed.change(on_change_seed_value, seed)
+ seed_checkbox.change(on_change_seed_checkbox, seed_checkbox)
+ safety_checker_checkbox.change(
+ on_change_safety_checker_checkbox, safety_checker_checkbox
+ )
+ tiny_auto_encoder_checkbox.change(
+ on_change_tiny_auto_encoder_checkbox, tiny_auto_encoder_checkbox
+ )
+ offline_checkbox.change(on_offline_checkbox, offline_checkbox)
+ img_format.change(on_change_image_format, img_format)
diff --git a/frontend/webui/image_to_image_ui.py b/frontend/webui/image_to_image_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..582355658ccc81c3ef8442f8a7f8e346bf0d415f
--- /dev/null
+++ b/frontend/webui/image_to_image_ui.py
@@ -0,0 +1,129 @@
+from typing import Any
+import gradio as gr
+from backend.models.lcmdiffusion_setting import DiffusionTask
+from models.interface_types import InterfaceType
+from frontend.utils import is_reshape_required
+from constants import DEVICE
+from state import get_settings, get_context
+from concurrent.futures import ThreadPoolExecutor
+
+
+app_settings = get_settings()
+
+previous_width = 0
+previous_height = 0
+previous_model_id = ""
+previous_num_of_images = 0
+
+
+def generate_image_to_image(
+ prompt,
+ negative_prompt,
+ init_image,
+ strength,
+) -> Any:
+ context = get_context(InterfaceType.WEBUI)
+ global \
+ previous_height, \
+ previous_width, \
+ previous_model_id, \
+ previous_num_of_images, \
+ app_settings
+
+ app_settings.settings.lcm_diffusion_setting.prompt = prompt
+ app_settings.settings.lcm_diffusion_setting.negative_prompt = negative_prompt
+ app_settings.settings.lcm_diffusion_setting.init_image = init_image
+ app_settings.settings.lcm_diffusion_setting.strength = strength
+
+ app_settings.settings.lcm_diffusion_setting.diffusion_task = (
+ DiffusionTask.image_to_image.value
+ )
+ model_id = app_settings.settings.lcm_diffusion_setting.openvino_lcm_model_id
+ reshape = False
+ image_width = app_settings.settings.lcm_diffusion_setting.image_width
+ image_height = app_settings.settings.lcm_diffusion_setting.image_height
+ num_images = app_settings.settings.lcm_diffusion_setting.number_of_images
+ if app_settings.settings.lcm_diffusion_setting.use_openvino:
+ reshape = is_reshape_required(
+ previous_width,
+ image_width,
+ previous_height,
+ image_height,
+ previous_model_id,
+ model_id,
+ previous_num_of_images,
+ num_images,
+ )
+
+ with ThreadPoolExecutor(max_workers=1) as executor:
+ future = executor.submit(
+ context.generate_text_to_image,
+ app_settings.settings,
+ reshape,
+ DEVICE,
+ )
+ images = future.result()
+ context.save_images(
+ images,
+ app_settings.settings,
+ )
+
+ previous_width = image_width
+ previous_height = image_height
+ previous_model_id = model_id
+ previous_num_of_images = num_images
+ return images
+
+
+def get_image_to_image_ui() -> None:
+ with gr.Blocks():
+ with gr.Row():
+ with gr.Column():
+ input_image = gr.Image(label="Init image", type="pil")
+ with gr.Row():
+ prompt = gr.Textbox(
+ show_label=False,
+ lines=3,
+ placeholder="A fantasy landscape",
+ container=False,
+ )
+
+ generate_btn = gr.Button(
+ "Generate",
+ elem_id="generate_button",
+ scale=0,
+ )
+ negative_prompt = gr.Textbox(
+ label="Negative prompt (Works in LCM-LoRA mode, set guidance > 1.0):",
+ lines=1,
+ placeholder="",
+ )
+ strength = gr.Slider(
+ 0.1,
+ 1,
+ value=app_settings.settings.lcm_diffusion_setting.strength,
+ step=0.01,
+ label="Strength",
+ )
+
+ input_params = [
+ prompt,
+ negative_prompt,
+ input_image,
+ strength,
+ ]
+
+ with gr.Column():
+ output = gr.Gallery(
+ label="Generated images",
+ show_label=True,
+ elem_id="gallery",
+ columns=2,
+ height=512,
+ )
+
+ generate_btn.click(
+ fn=generate_image_to_image,
+ inputs=input_params,
+ outputs=output,
+ )
diff --git a/frontend/webui/image_variations_ui.py b/frontend/webui/image_variations_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..ce88605cfdc4a9abb0f775fff7b097e49501067f
--- /dev/null
+++ b/frontend/webui/image_variations_ui.py
@@ -0,0 +1,111 @@
+from typing import Any
+import gradio as gr
+from backend.models.lcmdiffusion_setting import DiffusionTask
+from models.interface_types import InterfaceType
+from frontend.utils import is_reshape_required
+from constants import DEVICE
+from state import get_settings, get_context
+from concurrent.futures import ThreadPoolExecutor
+
+app_settings = get_settings()
+
+
+previous_width = 0
+previous_height = 0
+previous_model_id = ""
+previous_num_of_images = 0
+
+
+def generate_image_variations(
+ init_image,
+ variation_strength,
+) -> Any:
+ context = get_context(InterfaceType.WEBUI)
+ global \
+ previous_height, \
+ previous_width, \
+ previous_model_id, \
+ previous_num_of_images, \
+ app_settings
+
+ app_settings.settings.lcm_diffusion_setting.init_image = init_image
+ app_settings.settings.lcm_diffusion_setting.strength = variation_strength
+ app_settings.settings.lcm_diffusion_setting.prompt = ""
+ app_settings.settings.lcm_diffusion_setting.negative_prompt = ""
+
+ app_settings.settings.lcm_diffusion_setting.diffusion_task = (
+ DiffusionTask.image_to_image.value
+ )
+ model_id = app_settings.settings.lcm_diffusion_setting.openvino_lcm_model_id
+ reshape = False
+ image_width = app_settings.settings.lcm_diffusion_setting.image_width
+ image_height = app_settings.settings.lcm_diffusion_setting.image_height
+ num_images = app_settings.settings.lcm_diffusion_setting.number_of_images
+ if app_settings.settings.lcm_diffusion_setting.use_openvino:
+ reshape = is_reshape_required(
+ previous_width,
+ image_width,
+ previous_height,
+ image_height,
+ previous_model_id,
+ model_id,
+ previous_num_of_images,
+ num_images,
+ )
+
+ with ThreadPoolExecutor(max_workers=1) as executor:
+ future = executor.submit(
+ context.generate_text_to_image,
+ app_settings.settings,
+ reshape,
+ DEVICE,
+ )
+ images = future.result()
+ context.save_images(images, app_settings.settings)
+
+ previous_width = image_width
+ previous_height = image_height
+ previous_model_id = model_id
+ previous_num_of_images = num_images
+ return images
+
+
+def get_image_variations_ui() -> None:
+ with gr.Blocks():
+ with gr.Row():
+ with gr.Column():
+ input_image = gr.Image(label="Init image", type="pil")
+ with gr.Row():
+ generate_btn = gr.Button(
+ "Generate",
+ elem_id="generate_button",
+ scale=0,
+ )
+
+ variation_strength = gr.Slider(
+ 0.1,
+ 1,
+ value=0.4,
+ step=0.01,
+ label="Variations Strength",
+ )
+
+ input_params = [
+ input_image,
+ variation_strength,
+ ]
+
+ with gr.Column():
+ output = gr.Gallery(
+ label="Generated images",
+ show_label=True,
+ elem_id="gallery",
+ columns=2,
+ height=512,
+ )
+
+ generate_btn.click(
+ fn=generate_image_variations,
+ inputs=input_params,
+ outputs=output,
+ )
diff --git a/frontend/webui/lora_models_ui.py b/frontend/webui/lora_models_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..bc7838d057cb68621e43d2fe3027fef00dc3285a
--- /dev/null
+++ b/frontend/webui/lora_models_ui.py
@@ -0,0 +1,185 @@
+import gradio as gr
+from os import path
+from backend.lora import (
+ get_lora_models,
+ get_active_lora_weights,
+ update_lora_weights,
+ load_lora_weight,
+)
+from state import get_settings, get_context
+from frontend.utils import get_valid_lora_model
+from models.interface_types import InterfaceType
+
+
+_MAX_LORA_WEIGHTS = 5
+
+_custom_lora_sliders = []
+_custom_lora_names = []
+_custom_lora_columns = []
+
+app_settings = get_settings()
+
+
+def on_click_update_weight(*lora_weights):
+ update_weights = []
+ active_weights = get_active_lora_weights()
+ if not len(active_weights):
+ gr.Warning("No active LoRAs, first you need to load LoRA model")
+ return
+ for idx, lora in enumerate(active_weights):
+ update_weights.append(
+ (
+ lora[0],
+ lora_weights[idx],
+ )
+ )
+ if len(update_weights) > 0:
+ update_lora_weights(
+ get_context(InterfaceType.WEBUI).lcm_text_to_image.pipeline,
+ app_settings.settings.lcm_diffusion_setting,
+ update_weights,
+ )
+
+
+def on_click_load_lora(lora_name, lora_weight):
+ if app_settings.settings.lcm_diffusion_setting.use_openvino:
+ gr.Warning("Currently LoRA is not supported in OpenVINO.")
+ return
+ lora_models_map = get_lora_models(
+ app_settings.settings.lcm_diffusion_setting.lora.models_dir
+ )
+
+ # Load a new LoRA
+ settings = app_settings.settings.lcm_diffusion_setting
+ settings.lora.fuse = False
+ settings.lora.enabled = False
+ print(f"Selected Lora Model :{lora_name}")
+ print(f"Lora weight :{lora_weight}")
+ settings.lora.path = lora_models_map[lora_name]
+ settings.lora.weight = lora_weight
+ if not path.exists(settings.lora.path):
+ gr.Warning("Invalid LoRA model path!")
+ return
+ pipeline = get_context(InterfaceType.WEBUI).lcm_text_to_image.pipeline
+ if not pipeline:
+ gr.Warning("Pipeline not initialized. Please generate an image first.")
+ return
+ settings.lora.enabled = True
+ load_lora_weight(
+ get_context(InterfaceType.WEBUI).lcm_text_to_image.pipeline,
+ settings,
+ )
+
+ # Update Gradio LoRA UI
+ global _MAX_LORA_WEIGHTS
+ values = []
+ labels = []
+ rows = []
+ active_weights = get_active_lora_weights()
+ for idx, lora in enumerate(active_weights):
+ labels.append(f"{lora[0]}: ")
+ values.append(lora[1])
+ rows.append(gr.Row.update(visible=True))
+ for i in range(len(active_weights), _MAX_LORA_WEIGHTS):
+ labels.append(f"Update weight")
+ values.append(0.0)
+ rows.append(gr.Row.update(visible=False))
+ return labels + values + rows
+
+
+def get_lora_models_ui() -> None:
+ with gr.Blocks() as ui:
+ gr.HTML(
+ "Download and place your LoRA model weights in lora_models folders and restart App"
+ )
+ with gr.Row():
+ with gr.Column():
+ with gr.Row():
+ lora_models_map = get_lora_models(
+ app_settings.settings.lcm_diffusion_setting.lora.models_dir
+ )
+ valid_model = get_valid_lora_model(
+ list(lora_models_map.values()),
+ app_settings.settings.lcm_diffusion_setting.lora.path,
+ app_settings.settings.lcm_diffusion_setting.lora.models_dir,
+ )
+ if valid_model != "":
+ valid_model_path = lora_models_map[valid_model]
+ app_settings.settings.lcm_diffusion_setting.lora.path = (
+ valid_model_path
+ )
+ else:
+ app_settings.settings.lcm_diffusion_setting.lora.path = ""
+
+ lora_model = gr.Dropdown(
+ lora_models_map.keys(),
+ label="LoRA model",
+ info="LoRA model weight to load (You can use Lora models from Civitai or Hugging Face .safetensors format)",
+ value=valid_model,
+ interactive=True,
+ )
+
+ lora_weight = gr.Slider(
+ 0.0,
+ 1.0,
+ value=app_settings.settings.lcm_diffusion_setting.lora.weight,
+ step=0.05,
+ label="Initial Lora weight",
+ interactive=True,
+ )
+ load_lora_btn = gr.Button(
+ "Load selected LoRA",
+ elem_id="load_lora_button",
+ scale=0,
+ )
+
+ with gr.Row():
+ gr.Markdown(
+ "## Loaded LoRA models",
+ show_label=False,
+ )
+ update_lora_weights_btn = gr.Button(
+ "Update LoRA weights",
+ elem_id="load_lora_button",
+ scale=0,
+ )
+
+ global _MAX_LORA_WEIGHTS
+ global _custom_lora_sliders
+ global _custom_lora_names
+ global _custom_lora_columns
+ for i in range(0, _MAX_LORA_WEIGHTS):
+ new_row = gr.Column(visible=False)
+ _custom_lora_columns.append(new_row)
+ with new_row:
+ lora_name = gr.Markdown(
+ "Lora Name",
+ show_label=True,
+ )
+ lora_slider = gr.Slider(
+ 0.0,
+ 1.0,
+ step=0.05,
+ label="LoRA weight",
+ interactive=True,
+ visible=True,
+ )
+
+ _custom_lora_names.append(lora_name)
+ _custom_lora_sliders.append(lora_slider)
+
+ load_lora_btn.click(
+ fn=on_click_load_lora,
+ inputs=[lora_model, lora_weight],
+ outputs=[
+ *_custom_lora_names,
+ *_custom_lora_sliders,
+ *_custom_lora_columns,
+ ],
+ )
+
+ update_lora_weights_btn.click(
+ fn=on_click_update_weight,
+ inputs=[*_custom_lora_sliders],
+ outputs=None,
+ )
diff --git a/frontend/webui/models_ui.py b/frontend/webui/models_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..e121879baa1735fb9faec6a595a8dc59dd3a962b
--- /dev/null
+++ b/frontend/webui/models_ui.py
@@ -0,0 +1,180 @@
+import gradio as gr
+from constants import LCM_DEFAULT_MODEL
+from state import get_settings
+from frontend.utils import get_valid_model_id
+
+app_settings = get_settings()
+app_settings.settings.lcm_diffusion_setting.openvino_lcm_model_id = get_valid_model_id(
+ app_settings.openvino_lcm_models,
+ app_settings.settings.lcm_diffusion_setting.openvino_lcm_model_id,
+)
+
+
+def change_lcm_model_id(model_id):
+ app_settings.settings.lcm_diffusion_setting.lcm_model_id = model_id
+
+
+def change_lcm_lora_model_id(model_id):
+ app_settings.settings.lcm_diffusion_setting.lcm_lora.lcm_lora_id = model_id
+
+
+def change_lcm_lora_base_model_id(model_id):
+ app_settings.settings.lcm_diffusion_setting.lcm_lora.base_model_id = model_id
+
+
+def change_openvino_lcm_model_id(model_id):
+ app_settings.settings.lcm_diffusion_setting.openvino_lcm_model_id = model_id
+
+
+def change_gguf_diffusion_model(model_path):
+ if model_path == "None":
+ app_settings.settings.lcm_diffusion_setting.gguf_model.diffusion_path = ""
+ else:
+ app_settings.settings.lcm_diffusion_setting.gguf_model.diffusion_path = (
+ model_path
+ )
+
+
+def change_gguf_clip_model(model_path):
+ if model_path == "None":
+ app_settings.settings.lcm_diffusion_setting.gguf_model.clip_path = ""
+ else:
+ app_settings.settings.lcm_diffusion_setting.gguf_model.clip_path = model_path
+
+
+def change_gguf_t5xxl_model(model_path):
+ if model_path == "None":
+ app_settings.settings.lcm_diffusion_setting.gguf_model.t5xxl_path = ""
+ else:
+ app_settings.settings.lcm_diffusion_setting.gguf_model.t5xxl_path = model_path
+
+
+def change_gguf_vae_model(model_path):
+ if model_path == "None":
+ app_settings.settings.lcm_diffusion_setting.gguf_model.vae_path = ""
+ else:
+ app_settings.settings.lcm_diffusion_setting.gguf_model.vae_path = model_path
+
+
+def get_models_ui() -> None:
+ with gr.Blocks():
+ with gr.Row():
+ lcm_model_id = gr.Dropdown(
+ app_settings.lcm_models,
+ label="LCM model",
+ info="Diffusers LCM model ID",
+ value=get_valid_model_id(
+ app_settings.lcm_models,
+ app_settings.settings.lcm_diffusion_setting.lcm_model_id,
+ LCM_DEFAULT_MODEL,
+ ),
+ interactive=True,
+ )
+ with gr.Row():
+ lcm_lora_model_id = gr.Dropdown(
+ app_settings.lcm_lora_models,
+ label="LCM LoRA model",
+ info="Diffusers LCM LoRA model ID",
+ value=get_valid_model_id(
+ app_settings.lcm_lora_models,
+ app_settings.settings.lcm_diffusion_setting.lcm_lora.lcm_lora_id,
+ ),
+ interactive=True,
+ )
+ lcm_lora_base_model_id = gr.Dropdown(
+ app_settings.stable_diffsuion_models,
+ label="LCM LoRA base model",
+ info="Diffusers LCM LoRA base model ID",
+ value=get_valid_model_id(
+ app_settings.stable_diffsuion_models,
+ app_settings.settings.lcm_diffusion_setting.lcm_lora.base_model_id,
+ ),
+ interactive=True,
+ )
+ with gr.Row():
+ lcm_openvino_model_id = gr.Dropdown(
+ app_settings.openvino_lcm_models,
+ label="LCM OpenVINO model",
+ info="OpenVINO LCM-LoRA fused model ID",
+ value=get_valid_model_id(
+ app_settings.openvino_lcm_models,
+ app_settings.settings.lcm_diffusion_setting.openvino_lcm_model_id,
+ ),
+ interactive=True,
+ )
+ with gr.Row():
+ gguf_diffusion_model_id = gr.Dropdown(
+ app_settings.gguf_diffusion_models,
+ label="GGUF diffusion model",
+ info="GGUF diffusion model ",
+ value=get_valid_model_id(
+ app_settings.gguf_diffusion_models,
+ app_settings.settings.lcm_diffusion_setting.gguf_model.diffusion_path,
+ ),
+ interactive=True,
+ )
+ with gr.Row():
+ gguf_clip_model_id = gr.Dropdown(
+ app_settings.gguf_clip_models,
+ label="GGUF CLIP model",
+ info="GGUF CLIP model ",
+ value=get_valid_model_id(
+ app_settings.gguf_clip_models,
+ app_settings.settings.lcm_diffusion_setting.gguf_model.clip_path,
+ ),
+ interactive=True,
+ )
+ gguf_t5xxl_model_id = gr.Dropdown(
+ app_settings.gguf_t5xxl_models,
+ label="GGUF T5-XXL model",
+ info="GGUF T5-XXL model ",
+ value=get_valid_model_id(
+ app_settings.gguf_t5xxl_models,
+ app_settings.settings.lcm_diffusion_setting.gguf_model.t5xxl_path,
+ ),
+ interactive=True,
+ )
+ with gr.Row():
+ gguf_vae_model_id = gr.Dropdown(
+ app_settings.gguf_vae_models,
+ label="GGUF VAE model",
+ info="GGUF VAE model ",
+ value=get_valid_model_id(
+ app_settings.gguf_vae_models,
+ app_settings.settings.lcm_diffusion_setting.gguf_model.vae_path,
+ ),
+ interactive=True,
+ )
+
+ lcm_model_id.change(
+ change_lcm_model_id,
+ lcm_model_id,
+ )
+ lcm_lora_model_id.change(
+ change_lcm_lora_model_id,
+ lcm_lora_model_id,
+ )
+ lcm_lora_base_model_id.change(
+ change_lcm_lora_base_model_id,
+ lcm_lora_base_model_id,
+ )
+ lcm_openvino_model_id.change(
+ change_openvino_lcm_model_id,
+ lcm_openvino_model_id,
+ )
+ gguf_diffusion_model_id.change(
+ change_gguf_diffusion_model,
+ gguf_diffusion_model_id,
+ )
+ gguf_clip_model_id.change(
+ change_gguf_clip_model,
+ gguf_clip_model_id,
+ )
+ gguf_t5xxl_model_id.change(
+ change_gguf_t5xxl_model,
+ gguf_t5xxl_model_id,
+ )
+ gguf_vae_model_id.change(
+ change_gguf_vae_model,
+ gguf_vae_model_id,
+ )
diff --git a/frontend/webui/realtime_ui.py b/frontend/webui/realtime_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..0c556adcd674052bdec3ff6b2a9d3fb18c84a410
--- /dev/null
+++ b/frontend/webui/realtime_ui.py
@@ -0,0 +1,147 @@
+import base64
+from datetime import datetime
+from time import perf_counter
+
+import gradio as gr
+import numpy as np
+from backend.device import get_device_name, is_openvino_device
+from backend.lcm_text_to_image import LCMTextToImage
+from backend.models.lcmdiffusion_setting import LCMDiffusionSetting, LCMLora
+from constants import APP_VERSION, DEVICE
+from cv2 import imencode
+
+lcm_text_to_image = LCMTextToImage()
+lcm_lora = LCMLora(
+ base_model_id="Lykon/dreamshaper-8",
+ lcm_lora_id="latent-consistency/lcm-lora-sdv1-5",
+)
+
+
+# https://github.com/gradio-app/gradio/issues/2635#issuecomment-1423531319
+def encode_pil_to_base64_new(pil_image):
+ image_arr = np.asarray(pil_image)[:, :, ::-1]
+ _, byte_data = imencode(".png", image_arr)
+ base64_data = base64.b64encode(byte_data)
+ base64_string_opencv = base64_data.decode("utf-8")
+ return "data:image/png;base64," + base64_string_opencv
+
+
+# monkey patching encode pil
+gr.processing_utils.encode_pil_to_base64 = encode_pil_to_base64_new
+
+
+def predict(
+ prompt,
+ steps,
+ seed,
+):
+ lcm_diffusion_setting = LCMDiffusionSetting()
+ lcm_diffusion_setting.openvino_lcm_model_id = "rupeshs/sdxs-512-0.9-openvino"
+ lcm_diffusion_setting.prompt = prompt
+ lcm_diffusion_setting.guidance_scale = 1.0
+ lcm_diffusion_setting.inference_steps = steps
+ lcm_diffusion_setting.seed = seed
+ lcm_diffusion_setting.use_seed = True
+ lcm_diffusion_setting.image_width = 512
+ lcm_diffusion_setting.image_height = 512
+ lcm_diffusion_setting.use_openvino = True if is_openvino_device() else False
+ lcm_diffusion_setting.use_tiny_auto_encoder = True
+ lcm_text_to_image.init(
+ DEVICE,
+ lcm_diffusion_setting,
+ )
+ start = perf_counter()
+
+ images = lcm_text_to_image.generate(lcm_diffusion_setting)
+ latency = perf_counter() - start
+ print(f"Latency: {latency:.2f} seconds")
+ return images[0]
+
+
+css = """
+#container{
+ margin: 0 auto;
+ max-width: 40rem;
+}
+#intro{
+ max-width: 100%;
+ text-align: center;
+ margin: 0 auto;
+}
+#generate_button {
+ color: white;
+ border-color: #007bff;
+ background: #007bff;
+ width: 200px;
+ height: 50px;
+}
+footer {
+ visibility: hidden
+}
+"""
+
+
+def _get_footer_message() -> str:
+ version = f" {APP_VERSION} "
+ current_year = datetime.now().year
+ footer_msg = version + (
+ f' © 2023 - {current_year} '
+ " Rupesh Sreeraman
"
+ )
+ return footer_msg
+
+
+with gr.Blocks(css=css) as demo:
+ with gr.Column(elem_id="container"):
+ use_openvino = "- OpenVINO" if is_openvino_device() else ""
+ gr.Markdown(
+ f"""# Realtime FastSD CPU {use_openvino}
+ **Device : {DEVICE} , {get_device_name()}**
+ """,
+ elem_id="intro",
+ )
+
+ with gr.Row():
+ with gr.Row():
+ prompt = gr.Textbox(
+ placeholder="Describe the image you'd like to see",
+ scale=5,
+ container=False,
+ )
+ generate_btn = gr.Button(
+ "Generate",
+ scale=1,
+ elem_id="generate_button",
+ )
+
+ image = gr.Image(type="filepath")
+
+ steps = gr.Slider(
+ label="Steps",
+ value=1,
+ minimum=1,
+ maximum=6,
+ step=1,
+ visible=False,
+ )
+ seed = gr.Slider(
+ randomize=True,
+ minimum=0,
+ maximum=999999999,
+ label="Seed",
+ step=1,
+ )
+ gr.HTML(_get_footer_message())
+
+ inputs = [prompt, steps, seed]
+ prompt.input(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+ generate_btn.click(
+ fn=predict, inputs=inputs, outputs=image, show_progress=False
+ )
+ steps.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+ seed.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+
+
+def start_realtime_text_to_image(share=False):
+ demo.queue()
+ demo.launch(share=share)
diff --git a/frontend/webui/text_to_image_ui.py b/frontend/webui/text_to_image_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..0126946e9d2fe8ecfd6584c3911dca250848aa6c
--- /dev/null
+++ b/frontend/webui/text_to_image_ui.py
@@ -0,0 +1,105 @@
+import gradio as gr
+from typing import Any
+from backend.models.lcmdiffusion_setting import DiffusionTask
+from models.interface_types import InterfaceType
+from constants import DEVICE
+from state import get_settings, get_context
+from frontend.utils import is_reshape_required
+from concurrent.futures import ThreadPoolExecutor
+
+app_settings = get_settings()
+
+previous_width = 0
+previous_height = 0
+previous_model_id = ""
+previous_num_of_images = 0
+
+
+def generate_text_to_image(
+ prompt,
+ neg_prompt,
+) -> Any:
+ context = get_context(InterfaceType.WEBUI)
+ global \
+ previous_height, \
+ previous_width, \
+ previous_model_id, \
+ previous_num_of_images, \
+ app_settings
+ app_settings.settings.lcm_diffusion_setting.prompt = prompt
+ app_settings.settings.lcm_diffusion_setting.negative_prompt = neg_prompt
+ app_settings.settings.lcm_diffusion_setting.diffusion_task = (
+ DiffusionTask.text_to_image.value
+ )
+ model_id = app_settings.settings.lcm_diffusion_setting.openvino_lcm_model_id
+ reshape = False
+ image_width = app_settings.settings.lcm_diffusion_setting.image_width
+ image_height = app_settings.settings.lcm_diffusion_setting.image_height
+ num_images = app_settings.settings.lcm_diffusion_setting.number_of_images
+ if app_settings.settings.lcm_diffusion_setting.use_openvino:
+ reshape = is_reshape_required(
+ previous_width,
+ image_width,
+ previous_height,
+ image_height,
+ previous_model_id,
+ model_id,
+ previous_num_of_images,
+ num_images,
+ )
+
+ with ThreadPoolExecutor(max_workers=1) as executor:
+ future = executor.submit(
+ context.generate_text_to_image,
+ app_settings.settings,
+ reshape,
+ DEVICE,
+ )
+ images = future.result()
+ context.save_images(images, app_settings.settings)
+
+ previous_width = image_width
+ previous_height = image_height
+ previous_model_id = model_id
+ previous_num_of_images = num_images
+ return images
+
+
+def get_text_to_image_ui() -> None:
+ with gr.Blocks():
+ with gr.Row():
+ with gr.Column():
+ with gr.Row():
+ prompt = gr.Textbox(
+ show_label=False,
+ lines=3,
+ placeholder="A fantasy landscape",
+ container=False,
+ )
+
+ generate_btn = gr.Button(
+ "Generate",
+ elem_id="generate_button",
+ scale=0,
+ )
+ negative_prompt = gr.Textbox(
+ label="Negative prompt (Works in LCM-LoRA mode, set guidance > 1.0) :",
+ lines=1,
+ placeholder="",
+ )
+
+ input_params = [prompt, negative_prompt]
+
+ with gr.Column():
+ output = gr.Gallery(
+ label="Generated images",
+ show_label=True,
+ elem_id="gallery",
+ columns=2,
+ height=512,
+ )
+ generate_btn.click(
+ fn=generate_text_to_image,
+ inputs=input_params,
+ outputs=output,
+ )
diff --git a/frontend/webui/ui.py b/frontend/webui/ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..25c1dce5575cefcf07e3e1279054f10d35d99ed9
--- /dev/null
+++ b/frontend/webui/ui.py
@@ -0,0 +1,105 @@
+from datetime import datetime
+
+import gradio as gr
+from backend.device import get_device_name
+from constants import APP_VERSION
+from frontend.webui.controlnet_ui import get_controlnet_ui
+from frontend.webui.generation_settings_ui import get_generation_settings_ui
+from frontend.webui.image_to_image_ui import get_image_to_image_ui
+from frontend.webui.image_variations_ui import get_image_variations_ui
+from frontend.webui.lora_models_ui import get_lora_models_ui
+from frontend.webui.models_ui import get_models_ui
+from frontend.webui.text_to_image_ui import get_text_to_image_ui
+from frontend.webui.upscaler_ui import get_upscaler_ui
+from state import get_settings
+
+app_settings = get_settings()
+
+
+def _get_footer_message() -> str:
+ version = f" {APP_VERSION} "
+ current_year = datetime.now().year
+ footer_msg = version + (
+ f' © 2023 - {current_year} '
+ " Rupesh Sreeraman
"
+ )
+ return footer_msg
+
+
+def get_web_ui() -> gr.Blocks:
+ def change_mode(mode):
+ global app_settings
+ app_settings.settings.lcm_diffusion_setting.use_lcm_lora = False
+ app_settings.settings.lcm_diffusion_setting.use_openvino = False
+ app_settings.settings.lcm_diffusion_setting.use_gguf_model = False
+ if mode == "LCM-LoRA":
+ app_settings.settings.lcm_diffusion_setting.use_lcm_lora = True
+ elif mode == "LCM-OpenVINO":
+ app_settings.settings.lcm_diffusion_setting.use_openvino = True
+ elif mode == "GGUF":
+ app_settings.settings.lcm_diffusion_setting.use_gguf_model = True
+
+ # Prevent saved LoRA and ControlNet settings from being used by
+ # default; in WebUI mode, the user must explicitly enable those
+ if app_settings.settings.lcm_diffusion_setting.lora:
+ app_settings.settings.lcm_diffusion_setting.lora.enabled = False
+ if app_settings.settings.lcm_diffusion_setting.controlnet:
+ app_settings.settings.lcm_diffusion_setting.controlnet.enabled = False
+ theme = gr.themes.Default(
+ primary_hue="blue",
+ )
+ with gr.Blocks(
+ title="FastSD CPU",
+ theme=theme,
+ css="footer {visibility: hidden}",
+ ) as fastsd_web_ui:
+ gr.HTML("FastSD CPU
")
+ gr.Markdown(
+ f"**Processor : {get_device_name()}**",
+ elem_id="processor",
+ )
+ current_mode = "LCM"
+ if app_settings.settings.lcm_diffusion_setting.use_openvino:
+ current_mode = "LCM-OpenVINO"
+ elif app_settings.settings.lcm_diffusion_setting.use_lcm_lora:
+ current_mode = "LCM-LoRA"
+ elif app_settings.settings.lcm_diffusion_setting.use_gguf_model:
+ current_mode = "GGUF"
+
+ mode = gr.Radio(
+ ["LCM", "LCM-LoRA", "LCM-OpenVINO", "GGUF"],
+ label="Mode",
+ info="Current working mode",
+ value=current_mode,
+ )
+ mode.change(change_mode, inputs=mode)
+
+ with gr.Tabs():
+ with gr.TabItem("Text to Image"):
+ get_text_to_image_ui()
+ with gr.TabItem("Image to Image"):
+ get_image_to_image_ui()
+ with gr.TabItem("Image Variations"):
+ get_image_variations_ui()
+ with gr.TabItem("Upscaler"):
+ get_upscaler_ui()
+ with gr.TabItem("Generation Settings"):
+ get_generation_settings_ui()
+ with gr.TabItem("Models"):
+ get_models_ui()
+ with gr.TabItem("Lora Models"):
+ get_lora_models_ui()
+ with gr.TabItem("ControlNet"):
+ get_controlnet_ui()
+
+ gr.HTML(_get_footer_message())
+
+ return fastsd_web_ui
+
+
+def start_webui(
+ share: bool = False,
+):
+ webui = get_web_ui()
+ webui.queue()
+ webui.launch(share=share)
diff --git a/frontend/webui/upscaler_ui.py b/frontend/webui/upscaler_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..36a33b612f176315e1027fa137b69b668fd46610
--- /dev/null
+++ b/frontend/webui/upscaler_ui.py
@@ -0,0 +1,83 @@
+from typing import Any
+import gradio as gr
+from models.interface_types import InterfaceType
+from state import get_settings, get_context
+from backend.upscale.upscaler import upscale_image
+from backend.models.upscale import UpscaleMode
+from paths import FastStableDiffusionPaths
+
+app_settings = get_settings()
+
+
+previous_width = 0
+previous_height = 0
+previous_model_id = ""
+previous_num_of_images = 0
+
+
+def create_upscaled_image(
+ source_image,
+ upscale_mode,
+) -> Any:
+ context = get_context(InterfaceType.WEBUI)
+ scale_factor = 2
+ if upscale_mode == "SD":
+ mode = UpscaleMode.sd_upscale.value
+ elif upscale_mode == "AURA-SR":
+ mode = UpscaleMode.aura_sr.value
+ scale_factor = 4
+ else:
+ mode = UpscaleMode.normal.value
+
+ upscaled_filepath = FastStableDiffusionPaths.get_upscale_filepath(
+ None,
+ scale_factor,
+ app_settings.settings.generated_images.format,
+ )
+ image = upscale_image(
+ context=context,
+ src_image_path=source_image,
+ dst_image_path=upscaled_filepath,
+ upscale_mode=mode,
+ )
+ return image
+
+
+def get_upscaler_ui() -> None:
+ with gr.Blocks():
+ with gr.Row():
+ with gr.Column():
+ input_image = gr.Image(label="Image", type="filepath")
+ with gr.Row():
+ upscale_mode = gr.Radio(
+ ["EDSR", "SD", "AURA-SR"],
+ label="Upscale Mode (2x) | AURA-SR v2 (4x)",
+ info="Select upscale method(For SD Scale GGUF mode is not supported)",
+ value="EDSR",
+ )
+
+ generate_btn = gr.Button(
+ "Upscale",
+ elem_id="generate_button",
+ scale=0,
+ )
+
+ input_params = [
+ input_image,
+ upscale_mode,
+ ]
+
+ with gr.Column():
+ output = gr.Gallery(
+ label="Generated images",
+ show_label=True,
+ elem_id="gallery",
+ columns=2,
+ height=512,
+ )
+
+ generate_btn.click(
+ fn=create_upscaled_image,
+ inputs=input_params,
+ outputs=output,
+ )
diff --git a/image_ops.py b/image_ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..b60e911d37616bf29592b15fca9901f404a6e397
--- /dev/null
+++ b/image_ops.py
@@ -0,0 +1,15 @@
+from PIL import Image
+
+
+def resize_pil_image(
+ pil_image: Image,
+ image_width,
+ image_height,
+):
+ return pil_image.convert("RGB").resize(
+ (
+ image_width,
+ image_height,
+ ),
+ Image.Resampling.LANCZOS,
+ )
diff --git a/models/interface_types.py b/models/interface_types.py
new file mode 100644
index 0000000000000000000000000000000000000000..00f5c0a35fd3cc1f733edc2763e64781e2b7b6a9
--- /dev/null
+++ b/models/interface_types.py
@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class InterfaceType(Enum):
+ WEBUI = "Web User Interface"
+ GUI = "Graphical User Interface"
+ CLI = "Command Line Interface"
+ API_SERVER = "API Server"
diff --git a/models/settings.py b/models/settings.py
new file mode 100644
index 0000000000000000000000000000000000000000..afa1650fb5dab9dcef5645ac4aaf184939e8ae49
--- /dev/null
+++ b/models/settings.py
@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+from backend.models.lcmdiffusion_setting import LCMDiffusionSetting, LCMLora
+from backend.models.gen_images import GeneratedImages
+
+
+class Settings(BaseModel):
+ lcm_diffusion_setting: LCMDiffusionSetting = LCMDiffusionSetting(lcm_lora=LCMLora())
+ generated_images: GeneratedImages = GeneratedImages()
diff --git a/paths.py b/paths.py
new file mode 100644
index 0000000000000000000000000000000000000000..97f9af8b649e124f720560c63f20d1d8ad6f4778
--- /dev/null
+++ b/paths.py
@@ -0,0 +1,110 @@
+import os
+import constants
+from pathlib import Path
+from time import time
+from utils import get_image_file_extension
+
+
+def join_paths(
+ first_path: str,
+ second_path: str,
+) -> str:
+ return os.path.join(first_path, second_path)
+
+
+def get_file_name(file_path: str) -> str:
+ return Path(file_path).stem
+
+
+def get_app_path() -> str:
+ app_dir = os.path.dirname(__file__)
+ work_dir = os.path.dirname(app_dir)
+ return work_dir
+
+
+def get_configs_path() -> str:
+ config_path = join_paths(get_app_path(), constants.CONFIG_DIRECTORY)
+ return config_path
+
+
+class FastStableDiffusionPaths:
+ @staticmethod
+ def get_app_settings_path() -> str:
+ configs_path = get_configs_path()
+ settings_path = join_paths(
+ configs_path,
+ constants.APP_SETTINGS_FILE,
+ )
+ return settings_path
+
+ @staticmethod
+ def get_results_path() -> str:
+ results_path = join_paths(get_app_path(), constants.RESULTS_DIRECTORY)
+ return results_path
+
+ @staticmethod
+ def get_css_path() -> str:
+ app_dir = os.path.dirname(__file__)
+ css_path = os.path.join(
+ app_dir,
+ "frontend",
+ "webui",
+ "css",
+ "style.css",
+ )
+ return css_path
+
+ @staticmethod
+ def get_models_config_path(model_config_file: str) -> str:
+ configs_path = get_configs_path()
+ models_path = join_paths(
+ configs_path,
+ model_config_file,
+ )
+ return models_path
+
+ @staticmethod
+ def get_upscale_filepath(
+ file_path_src: str,
+ scale_factor: int,
+ format: str,
+ ) -> str:
+ if file_path_src:
+ file_name_src = get_file_name(file_path_src)
+ else:
+ file_name_src = "fastsdcpu"
+
+ extension = get_image_file_extension(format)
+ upscaled_filepath = join_paths(
+ FastStableDiffusionPaths.get_results_path(),
+ f"{file_name_src}_{int(scale_factor)}x_upscale_{int(time())}{extension}",
+ )
+ return upscaled_filepath
+
+ @staticmethod
+ def get_lora_models_path() -> str:
+ lora_models_path = join_paths(get_app_path(), constants.LORA_DIRECTORY)
+ return lora_models_path
+
+ @staticmethod
+ def get_controlnet_models_path() -> str:
+ controlnet_models_path = join_paths(
+ get_app_path(), constants.CONTROLNET_DIRECTORY
+ )
+ return controlnet_models_path
+
+ @staticmethod
+ def get_gguf_models_path() -> str:
+ models_path = join_paths(get_app_path(), constants.MODELS_DIRECTORY)
+ guuf_models_path = join_paths(models_path, "gguf")
+ return guuf_models_path
+
+
+def get_base_folder_name(path: str) -> str:
+ return os.path.basename(path)
+
+
+def ensure_path(path: str) -> None:
+ """Ensure that the directory exists."""
+ if not os.path.exists(path):
+ os.makedirs(path, exist_ok=True)
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..152c7ea5ce2e89b210693374be56e3c8cc5684f9
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,20 @@
+accelerate==0.33.0
+diffusers==0.30.0
+transformers==4.41.2
+PyQt5
+Pillow==9.4.0
+openvino==2024.4.0
+optimum-intel==1.18.2
+onnx==1.16.0
+onnxruntime==1.17.3
+pydantic
+typing-extensions==4.8.0
+pyyaml==6.0.1
+gradio==5.6.0
+peft==0.6.1
+opencv-python==4.8.1.78
+omegaconf==2.3.0
+controlnet-aux==0.0.7
+mediapipe>=0.10.9
+tomesd==0.1.3
+fastapi-mcp
\ No newline at end of file
diff --git a/state.py b/state.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c0c4d5599c192840646f8d655cc16e4fbb8c6b2
--- /dev/null
+++ b/state.py
@@ -0,0 +1,32 @@
+from app_settings import AppSettings
+from typing import Optional
+
+from context import Context
+from models.interface_types import InterfaceType
+
+
+class _AppState:
+ _instance: Optional["_AppState"] = None
+ settings: Optional[AppSettings] = None
+ context: Optional[Context] = None
+
+
+def get_state() -> _AppState:
+ if _AppState._instance is None:
+ _AppState._instance = _AppState()
+ return _AppState._instance
+
+
+def get_settings(skip_file: bool = False) -> AppSettings:
+ state = get_state()
+ if state.settings is None:
+ state.settings = AppSettings()
+ state.settings.load(skip_file)
+ return state.settings
+
+
+def get_context(interface_type: InterfaceType) -> Context:
+ state = get_state()
+ if state.context is None:
+ state.context = Context(interface_type)
+ return state.context
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..a6d2d9bcb0840905089c22f682694ea567233572
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,38 @@
+from os import path, listdir
+import platform
+from typing import List
+
+
+def show_system_info():
+ try:
+ print(f"Running on {platform.system()} platform")
+ print(f"OS: {platform.platform()}")
+ print(f"Processor: {platform.processor()}")
+ except Exception as ex:
+ print(f"Error occurred while getting system information {ex}")
+
+
+def get_models_from_text_file(file_path: str) -> List:
+ models = []
+ with open(file_path, "r") as file:
+ lines = file.readlines()
+ for repo_id in lines:
+ if repo_id.strip() != "":
+ models.append(repo_id.strip())
+ return models
+
+
+def get_image_file_extension(image_format: str) -> str:
+ if image_format == "JPEG":
+ return ".jpg"
+ elif image_format == "PNG":
+ return ".png"
+
+
+def get_files_in_dir(root_dir: str) -> List:
+ models = []
+ models.append("None")
+ for file in listdir(root_dir):
+ if file.endswith((".gguf", ".safetensors")):
+ models.append(path.join(root_dir, file))
+ return models