Merge pull request #105 from rupeshs/add-image-variations

Add image variations

Readme.md

@@ -94,6 +94,7 @@ You can directly use these models in FastSD CPU.
 - Fast 1 step inference (SDXL Turbo)
 - Added SD Turbo support
 - Added image to image support for Turbo models (Pytorch and OpenVINO)
+- Added image variations support
 
 ## 2 Steps fast inference
 FastSD CPU supports 2 to 3 steps fast inference using LCM-LoRA workflow. It works well with SD 1.5 models.

src/app.py

@@ -2,11 +2,11 @@
 from utils import show_system_info
 import constants
 from argparse import ArgumentParser
-from context import Context
+
 from constants import APP_VERSION, LCM_DEFAULT_MODEL_OPENVINO
 from models.interface_types import InterfaceType
 from constants import DEVICE
-from state import get_settings
+from state import get_settings, get_context
 
 parser = ArgumentParser(description=f"FAST SD CPU {constants.APP_VERSION}")
 parser.add_argument(
@@ -178,7 +178,7 @@
     print("Starting realtime text to image(EXPERIMENTAL)")
     start_realtime_text_to_image(args.share)
 else:
-    context = Context(InterfaceType.CLI)
+    context = get_context(InterfaceType.CLI)
     config = app_settings.settings
 
     if args.use_openvino:

src/backend/lcm_text_to_image.py

@@ -38,6 +38,7 @@ def __init__(
         self.previous_use_openvino = False
         self.img_to_img_pipeline = None
         self.is_openvino_init = False
+        self.task_type = DiffusionTask.text_to_image
         self.torch_data_type = (
             torch.float32 if is_openvino_device() or DEVICE == "mps" else torch.float16
         )
@@ -110,6 +111,7 @@ def init(
             or self.previous_ov_model_id != ov_model_id
             or self.previous_safety_checker != lcm_diffusion_setting.use_safety_checker
             or self.previous_use_openvino != lcm_diffusion_setting.use_openvino
+            or self.previous_task_type != lcm_diffusion_setting.diffusion_task
         ):
             if self.use_openvino and is_openvino_device():
                 if self.pipeline:
@@ -218,6 +220,7 @@ def init(
             self.previous_use_lcm_lora = use_lora
             self.previous_safety_checker = lcm_diffusion_setting.use_safety_checker
             self.previous_use_openvino = lcm_diffusion_setting.use_openvino
+            self.previous_task_type = lcm_diffusion_setting.diffusion_task
             if (
                 lcm_diffusion_setting.diffusion_task
                 == DiffusionTask.text_to_image.value

src/constants.py

@@ -1,6 +1,6 @@
 from os import environ
 
-APP_VERSION = "v1.0.0 beta 22"
+APP_VERSION = "v1.0.0 beta 23"
 LCM_DEFAULT_MODEL = "stabilityai/sd-turbo"
 LCM_DEFAULT_MODEL_OPENVINO = "rupeshs/sd-turbo-openvino"
 APP_NAME = "FastSD CPU"

src/context.py

@@ -5,7 +5,6 @@
 from time import perf_counter
 from backend.image_saver import ImageSaver
 from pprint import pprint
-from state import get_settings
 
 
 class Context:
@@ -23,8 +22,10 @@ def generate_text_to_image(
         reshape: bool = False,
         device: str = "cpu",
     ) -> Any:
-        get_settings().save()
         tick = perf_counter()
+        from state import get_settings
+
+        get_settings().save()
         pprint(settings.lcm_diffusion_setting.model_dump())
         if not settings.lcm_diffusion_setting.lcm_lora:
             return None

src/frontend/webui/generation_settings_ui.py

@@ -137,4 +137,4 @@ def get_generation_settings_ui() -> None:
         tiny_auto_encoder_checkbox.change(
             on_change_tiny_auto_encoder_checkbox, tiny_auto_encoder_checkbox
         )
-        offline_checkbox.change(on_change_tiny_auto_encoder_checkbox, offline_checkbox)
+        offline_checkbox.change(on_offline_checkbox, offline_checkbox)

src/frontend/webui/image_to_image_ui.py

@@ -1,16 +1,15 @@
 from typing import Any
 import gradio as gr
 from backend.models.lcmdiffusion_setting import DiffusionTask
-from context import Context
 from models.interface_types import InterfaceType
 from frontend.utils import is_reshape_required
 from constants import DEVICE
-from state import get_settings
+from state import get_settings, get_context
 from concurrent.futures import ThreadPoolExecutor
 
 app_settings = get_settings()
 
-context = Context(InterfaceType.WEBUI)
+context = get_context(InterfaceType.WEBUI)
 previous_width = 0
 previous_height = 0
 previous_model_id = ""

src/frontend/webui/image_variations_ui.py

@@ -0,0 +1,105 @@
+from typing import Any
+import gradio as gr
+from backend.models.lcmdiffusion_setting import DiffusionTask
+from context import Context
+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()
+
+context = get_context(InterfaceType.WEBUI)
+previous_width = 0
+previous_height = 0
+previous_model_id = ""
+previous_num_of_images = 0
+
+
+def generate_image_variations(
+    init_image,
+    variation_strength,
+) -> Any:
+    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()
+
+    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,
+    )

src/frontend/webui/text_to_image_ui.py

@@ -1,16 +1,15 @@
 import gradio as gr
 from typing import Any
 from backend.models.lcmdiffusion_setting import DiffusionTask
-from context import Context
 from models.interface_types import InterfaceType
 from constants import DEVICE
-from state import get_settings
+from state import get_settings, get_context
 from frontend.utils import is_reshape_required
 from concurrent.futures import ThreadPoolExecutor
 from pprint import pprint
 
 app_settings = get_settings()
-context = Context(InterfaceType.WEBUI)
+context = get_context(InterfaceType.WEBUI)
 previous_width = 0
 previous_height = 0
 previous_model_id = ""

src/frontend/webui/ui.py

@@ -4,6 +4,7 @@
 from frontend.webui.image_to_image_ui import get_image_to_image_ui
 from frontend.webui.generation_settings_ui import get_generation_settings_ui
 from frontend.webui.models_ui import get_models_ui
+from frontend.webui.image_variations_ui import get_image_variations_ui
 from paths import FastStableDiffusionPaths
 from state import get_settings
 
@@ -53,6 +54,8 @@ def change_mode(mode):
                 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("Generation Settings"):
                 get_generation_settings_ui()
             with gr.TabItem("Models"):

src/state.py

@@ -1,10 +1,14 @@
 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:
@@ -19,3 +23,10 @@ def get_settings(skip_file: bool = False) -> AppSettings:
         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