Upload 11 files

src/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 <a href="https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/tree/main">ControlNet v1.1 </a> (723 MB files) and place it in <b>controlnet_models</b> 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
+    )

src/frontend/webui/errors.py

@@ -0,0 +1,10 @@
+import gradio as gr
+
+
+def show_error(message: str) -> None:
+    if "num_inference_steps != 2" in message:
+        gr.Warning(
+            "Please set the generation setting inference steps to 2 for SANA sprint model"
+        )
+    else:
+        gr.Warning(message)

src/frontend/webui/generation_settings_ui.py

@@ -0,0 +1,183 @@
+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 AutoEncoder (Fast, moderate quality)",
+                    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)

src/frontend/webui/image_to_image_ui.py

@@ -0,0 +1,132 @@
+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
+from frontend.webui.errors import show_error
+
+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()
+        if images:
+            context.save_images(
+                images,
+                app_settings.settings,
+            )
+        else:
+            show_error(context.error)
+
+    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,
+    )

src/frontend/webui/image_variations_ui.py

@@ -0,0 +1,118 @@
+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
+from frontend.webui.errors import show_error
+
+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()
+        if images:
+            context.save_images(
+                images,
+                app_settings.settings,
+            )
+        else:
+            show_error(context.error)
+
+    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/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 <b>lora_models</b> 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,
+    )

src/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,
+        )

src/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"<center><p> {APP_VERSION} "
+    current_year = datetime.now().year
+    footer_msg = version + (
+        f'  © 2023 - {current_year} <a href="https://github.com/rupeshs">'
+        " Rupesh Sreeraman</a></p></center>"
+    )
+    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)

src/frontend/webui/text_to_image_ui.py

@@ -0,0 +1,112 @@
+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
+from frontend.webui.errors import show_error
+
+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()
+        if images:
+            context.save_images(
+                images,
+                app_settings.settings,
+            )
+        else:
+            show_error(context.error)
+
+    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,
+    )

src/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"<center><p> {APP_VERSION} "
+    current_year = datetime.now().year
+    footer_msg = version + (
+        f'  © 2023 - {current_year} <a href="https://github.com/rupeshs">'
+        " Rupesh Sreeraman</a></p></center>"
+    )
+    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("<center><H1>FastSD CPU</H1></center>")
+        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)

src/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,
+    )