Upload 2 files

src/frontend/cli_interactive.py

@@ -0,0 +1,661 @@
+from os import path
+from PIL import Image
+from typing import Any
+
+from constants import DEVICE
+from paths import FastStableDiffusionPaths
+from backend.upscale.upscaler import upscale_image
+from backend.upscale.tiled_upscale import generate_upscaled_image
+from frontend.webui.image_variations_ui import generate_image_variations
+from backend.lora import (
+    get_active_lora_weights,
+    update_lora_weights,
+    load_lora_weight,
+)
+from backend.models.lcmdiffusion_setting import (
+    DiffusionTask,
+    ControlNetSetting,
+)
+
+
+_batch_count = 1
+_edit_lora_settings = False
+
+
+def user_value(
+    value_type: type,
+    message: str,
+    default_value: Any,
+) -> Any:
+    try:
+        value = value_type(input(message))
+    except:
+        value = default_value
+    return value
+
+
+def interactive_mode(
+    config,
+    context,
+):
+    print("=============================================")
+    print("Welcome to FastSD CPU Interactive CLI")
+    print("=============================================")
+    while True:
+        print("> 1. Text to Image")
+        print("> 2. Image to Image")
+        print("> 3. Image Variations")
+        print("> 4. EDSR Upscale")
+        print("> 5. SD Upscale")
+        print("> 6. Edit default generation settings")
+        print("> 7. Edit LoRA settings")
+        print("> 8. Edit ControlNet settings")
+        print("> 9. Edit negative prompt")
+        print("> 10. Quit")
+        option = user_value(
+            int,
+            "Enter a Diffusion Task number (1): ",
+            1,
+        )
+        if option not in range(1, 11):
+            print("Wrong Diffusion Task number!")
+            exit()
+
+        if option == 1:
+            interactive_txt2img(
+                config,
+                context,
+            )
+        elif option == 2:
+            interactive_img2img(
+                config,
+                context,
+            )
+        elif option == 3:
+            interactive_variations(
+                config,
+                context,
+            )
+        elif option == 4:
+            interactive_edsr(
+                config,
+                context,
+            )
+        elif option == 5:
+            interactive_sdupscale(
+                config,
+                context,
+            )
+        elif option == 6:
+            interactive_settings(
+                config,
+                context,
+            )
+        elif option == 7:
+            interactive_lora(
+                config,
+                context,
+                True,
+            )
+        elif option == 8:
+            interactive_controlnet(
+                config,
+                context,
+                True,
+            )
+        elif option == 9:
+            interactive_negative(
+                config,
+                context,
+            )
+        elif option == 10:
+            exit()
+
+
+def interactive_negative(
+    config,
+    context,
+):
+    settings = config.lcm_diffusion_setting
+    print(f"Current negative prompt: '{settings.negative_prompt}'")
+    user_input = input("Write a negative prompt (set guidance > 1.0): ")
+    if user_input == "":
+        return
+    else:
+        settings.negative_prompt = user_input
+
+
+def interactive_controlnet(
+    config,
+    context,
+    menu_flag=False,
+):
+    """
+    @param menu_flag: Indicates whether this function was called from the main
+        interactive CLI menu; _True_ if called from the main menu, _False_ otherwise
+    """
+    settings = config.lcm_diffusion_setting
+    if not settings.controlnet:
+        settings.controlnet = ControlNetSetting()
+
+    current_enabled = settings.controlnet.enabled
+    current_adapter_path = settings.controlnet.adapter_path
+    current_conditioning_scale = settings.controlnet.conditioning_scale
+    current_control_image = settings.controlnet._control_image
+
+    option = input("Enable ControlNet? (y/N): ")
+    settings.controlnet.enabled = True if option.upper() == "Y" else False
+    if settings.controlnet.enabled:
+        option = input(
+            f"Enter ControlNet adapter path ({settings.controlnet.adapter_path}): "
+        )
+        if option != "":
+            settings.controlnet.adapter_path = option
+        settings.controlnet.conditioning_scale = user_value(
+            float,
+            f"Enter ControlNet conditioning scale ({settings.controlnet.conditioning_scale}): ",
+            settings.controlnet.conditioning_scale,
+        )
+        option = input(
+            f"Enter ControlNet control image path (Leave empty to reuse current): "
+        )
+        if option != "":
+            try:
+                new_image = Image.open(option)
+                settings.controlnet._control_image = new_image
+            except (AttributeError, FileNotFoundError) as e:
+                settings.controlnet._control_image = None
+        if (
+            not settings.controlnet.adapter_path
+            or not path.exists(settings.controlnet.adapter_path)
+            or not settings.controlnet._control_image
+        ):
+            print("Invalid ControlNet settings! Disabling ControlNet")
+            settings.controlnet.enabled = False
+
+    if (
+        settings.controlnet.enabled != current_enabled
+        or settings.controlnet.adapter_path != current_adapter_path
+    ):
+        settings.rebuild_pipeline = True
+
+
+def interactive_lora(
+    config,
+    context,
+    menu_flag=False,
+):
+    """
+    @param menu_flag: Indicates whether this function was called from the main
+        interactive CLI menu; _True_ if called from the main menu, _False_ otherwise
+    """
+    if context == None or context.lcm_text_to_image.pipeline == None:
+        print("Diffusion pipeline not initialized, please run a generation task first!")
+        return
+
+    print("> 1. Change LoRA weights")
+    print("> 2. Load new LoRA model")
+    option = user_value(
+        int,
+        "Enter a LoRA option (1): ",
+        1,
+    )
+    if option not in range(1, 3):
+        print("Wrong LoRA option!")
+        return
+
+    if option == 1:
+        update_weights = []
+        active_weights = get_active_lora_weights()
+        for lora in active_weights:
+            weight = user_value(
+                float,
+                f"Enter a new LoRA weight for {lora[0]} ({lora[1]}): ",
+                lora[1],
+            )
+            update_weights.append(
+                (
+                    lora[0],
+                    weight,
+                )
+            )
+        if len(update_weights) > 0:
+            update_lora_weights(
+                context.lcm_text_to_image.pipeline,
+                config.lcm_diffusion_setting,
+                update_weights,
+            )
+    elif option == 2:
+        # Load a new LoRA
+        settings = config.lcm_diffusion_setting
+        settings.lora.fuse = False
+        settings.lora.enabled = False
+        settings.lora.path = input("Enter LoRA model path: ")
+        settings.lora.weight = user_value(
+            float,
+            "Enter a LoRA weight (0.5): ",
+            0.5,
+        )
+        if not path.exists(settings.lora.path):
+            print("Invalid LoRA model path!")
+            return
+        settings.lora.enabled = True
+        load_lora_weight(context.lcm_text_to_image.pipeline, settings)
+
+    if menu_flag:
+        global _edit_lora_settings
+        _edit_lora_settings = False
+        option = input("Edit LoRA settings after every generation? (y/N): ")
+        if option.upper() == "Y":
+            _edit_lora_settings = True
+
+
+def interactive_settings(
+    config,
+    context,
+):
+    global _batch_count
+    settings = config.lcm_diffusion_setting
+    print("Enter generation settings (leave empty to use current value)")
+    print("> 1. Use LCM")
+    print("> 2. Use LCM-Lora")
+    print("> 3. Use OpenVINO")
+    option = user_value(
+        int,
+        "Select inference model option (1): ",
+        1,
+    )
+    if option not in range(1, 4):
+        print("Wrong inference model option! Falling back to defaults")
+        return
+
+    settings.use_lcm_lora = False
+    settings.use_openvino = False
+    if option == 1:
+        lcm_model_id = input(f"Enter LCM model ID ({settings.lcm_model_id}): ")
+        if lcm_model_id != "":
+            settings.lcm_model_id = lcm_model_id
+    elif option == 2:
+        settings.use_lcm_lora = True
+        lcm_lora_id = input(
+            f"Enter LCM-Lora model ID ({settings.lcm_lora.lcm_lora_id}): "
+        )
+        if lcm_lora_id != "":
+            settings.lcm_lora.lcm_lora_id = lcm_lora_id
+        base_model_id = input(
+            f"Enter Base model ID ({settings.lcm_lora.base_model_id}): "
+        )
+        if base_model_id != "":
+            settings.lcm_lora.base_model_id = base_model_id
+    elif option == 3:
+        settings.use_openvino = True
+        openvino_lcm_model_id = input(
+            f"Enter OpenVINO model ID ({settings.openvino_lcm_model_id}): "
+        )
+        if openvino_lcm_model_id != "":
+            settings.openvino_lcm_model_id = openvino_lcm_model_id
+
+    settings.use_offline_model = True
+    settings.use_tiny_auto_encoder = True
+    option = input("Work offline? (Y/n): ")
+    if option.upper() == "N":
+        settings.use_offline_model = False
+    option = input("Use Tiny Auto Encoder? (Y/n): ")
+    if option.upper() == "N":
+        settings.use_tiny_auto_encoder = False
+
+    settings.image_width = user_value(
+        int,
+        f"Image width ({settings.image_width}): ",
+        settings.image_width,
+    )
+    settings.image_height = user_value(
+        int,
+        f"Image height ({settings.image_height}): ",
+        settings.image_height,
+    )
+    settings.inference_steps = user_value(
+        int,
+        f"Inference steps ({settings.inference_steps}): ",
+        settings.inference_steps,
+    )
+    settings.guidance_scale = user_value(
+        float,
+        f"Guidance scale ({settings.guidance_scale}): ",
+        settings.guidance_scale,
+    )
+    settings.number_of_images = user_value(
+        int,
+        f"Number of images per batch ({settings.number_of_images}): ",
+        settings.number_of_images,
+    )
+    _batch_count = user_value(
+        int,
+        f"Batch count ({_batch_count}): ",
+        _batch_count,
+    )
+    # output_format = user_value(int, f"Output format (PNG)", 1)
+    print(config.lcm_diffusion_setting)
+
+
+def interactive_txt2img(
+    config,
+    context,
+):
+    global _batch_count
+    config.lcm_diffusion_setting.diffusion_task = DiffusionTask.text_to_image.value
+    user_input = input("Write a prompt (write 'exit' to quit): ")
+    while True:
+        if user_input == "exit":
+            return
+        elif user_input == "":
+            user_input = config.lcm_diffusion_setting.prompt
+        config.lcm_diffusion_setting.prompt = user_input
+        for _ in range(0, _batch_count):
+            images = context.generate_text_to_image(
+                settings=config,
+                device=DEVICE,
+            )
+            context.save_images(
+                images,
+                config,
+            )
+        if _edit_lora_settings:
+            interactive_lora(
+                config,
+                context,
+            )
+        user_input = input("Write a prompt: ")
+
+
+def interactive_img2img(
+    config,
+    context,
+):
+    global _batch_count
+    settings = config.lcm_diffusion_setting
+    settings.diffusion_task = DiffusionTask.image_to_image.value
+    steps = settings.inference_steps
+    source_path = input("Image path: ")
+    if source_path == "":
+        print("Error : You need to provide a file in img2img mode")
+        return
+    settings.strength = user_value(
+        float,
+        f"img2img strength ({settings.strength}): ",
+        settings.strength,
+    )
+    settings.inference_steps = int(steps / settings.strength + 1)
+    user_input = input("Write a prompt (write 'exit' to quit): ")
+    while True:
+        if user_input == "exit":
+            settings.inference_steps = steps
+            return
+        settings.init_image = Image.open(source_path)
+        settings.prompt = user_input
+        for _ in range(0, _batch_count):
+            images = context.generate_text_to_image(
+                settings=config,
+                device=DEVICE,
+            )
+            context.save_images(
+                images,
+                config,
+            )
+        new_path = input(f"Image path ({source_path}): ")
+        if new_path != "":
+            source_path = new_path
+        settings.strength = user_value(
+            float,
+            f"img2img strength ({settings.strength}): ",
+            settings.strength,
+        )
+        if _edit_lora_settings:
+            interactive_lora(
+                config,
+                context,
+            )
+        settings.inference_steps = int(steps / settings.strength + 1)
+        user_input = input("Write a prompt: ")
+
+
+def interactive_variations(
+    config,
+    context,
+):
+    global _batch_count
+    settings = config.lcm_diffusion_setting
+    settings.diffusion_task = DiffusionTask.image_to_image.value
+    steps = settings.inference_steps
+    source_path = input("Image path: ")
+    if source_path == "":
+        print("Error : You need to provide a file in Image variations mode")
+        return
+    settings.strength = user_value(
+        float,
+        f"Image variations strength ({settings.strength}): ",
+        settings.strength,
+    )
+    settings.inference_steps = int(steps / settings.strength + 1)
+    while True:
+        settings.init_image = Image.open(source_path)
+        settings.prompt = ""
+        for i in range(0, _batch_count):
+            generate_image_variations(
+                settings.init_image,
+                settings.strength,
+            )
+        if _edit_lora_settings:
+            interactive_lora(
+                config,
+                context,
+            )
+        user_input = input("Continue in Image variations mode? (Y/n): ")
+        if user_input.upper() == "N":
+            settings.inference_steps = steps
+            return
+        new_path = input(f"Image path ({source_path}): ")
+        if new_path != "":
+            source_path = new_path
+        settings.strength = user_value(
+            float,
+            f"Image variations strength ({settings.strength}): ",
+            settings.strength,
+        )
+        settings.inference_steps = int(steps / settings.strength + 1)
+
+
+def interactive_edsr(
+    config,
+    context,
+):
+    source_path = input("Image path: ")
+    if source_path == "":
+        print("Error : You need to provide a file in EDSR mode")
+        return
+    while True:
+        output_path = FastStableDiffusionPaths.get_upscale_filepath(
+            source_path,
+            2,
+            config.generated_images.format,
+        )
+        result = upscale_image(
+            context,
+            source_path,
+            output_path,
+            2,
+        )
+        user_input = input("Continue in EDSR upscale mode? (Y/n): ")
+        if user_input.upper() == "N":
+            return
+        new_path = input(f"Image path ({source_path}): ")
+        if new_path != "":
+            source_path = new_path
+
+
+def interactive_sdupscale_settings(config):
+    steps = config.lcm_diffusion_setting.inference_steps
+    custom_settings = {}
+    print("> 1. Upscale whole image")
+    print("> 2. Define custom tiles (advanced)")
+    option = user_value(
+        int,
+        "Select an SD Upscale option (1): ",
+        1,
+    )
+    if option not in range(1, 3):
+        print("Wrong SD Upscale option!")
+        return
+
+    # custom_settings["source_file"] = args.file
+    custom_settings["source_file"] = ""
+    new_path = input(f"Input image path ({custom_settings['source_file']}): ")
+    if new_path != "":
+        custom_settings["source_file"] = new_path
+    if custom_settings["source_file"] == "":
+        print("Error : You need to provide a file in SD Upscale mode")
+        return
+    custom_settings["target_file"] = None
+    if option == 2:
+        custom_settings["target_file"] = input("Image to patch: ")
+        if custom_settings["target_file"] == "":
+            print("No target file provided, upscaling whole input image instead!")
+            custom_settings["target_file"] = None
+            option = 1
+    custom_settings["output_format"] = config.generated_images.format
+    custom_settings["strength"] = user_value(
+        float,
+        f"SD Upscale strength ({config.lcm_diffusion_setting.strength}): ",
+        config.lcm_diffusion_setting.strength,
+    )
+    config.lcm_diffusion_setting.inference_steps = int(
+        steps / custom_settings["strength"] + 1
+    )
+    if option == 1:
+        custom_settings["scale_factor"] = user_value(
+            float,
+            f"Scale factor (2.0): ",
+            2.0,
+        )
+        custom_settings["tile_size"] = user_value(
+            int,
+            f"Split input image into tiles of the following size, in pixels (256): ",
+            256,
+        )
+        custom_settings["tile_overlap"] = user_value(
+            int,
+            f"Tile overlap, in pixels (16): ",
+            16,
+        )
+    elif option == 2:
+        custom_settings["scale_factor"] = user_value(
+            float,
+            "Input image to Image-to-patch scale_factor (2.0): ",
+            2.0,
+        )
+        custom_settings["tile_size"] = 256
+        custom_settings["tile_overlap"] = 16
+    custom_settings["prompt"] = input(
+        "Write a prompt describing the input image (optional): "
+    )
+    custom_settings["tiles"] = []
+    if option == 2:
+        add_tile = True
+        while add_tile:
+            print("=== Define custom SD Upscale tile ===")
+            tile_x = user_value(
+                int,
+                "Enter tile's X position: ",
+                0,
+            )
+            tile_y = user_value(
+                int,
+                "Enter tile's Y position: ",
+                0,
+            )
+            tile_w = user_value(
+                int,
+                "Enter tile's width (256): ",
+                256,
+            )
+            tile_h = user_value(
+                int,
+                "Enter tile's height (256): ",
+                256,
+            )
+            tile_scale = user_value(
+                float,
+                "Enter tile's scale factor (2.0): ",
+                2.0,
+            )
+            tile_prompt = input("Enter tile's prompt (optional): ")
+            custom_settings["tiles"].append(
+                {
+                    "x": tile_x,
+                    "y": tile_y,
+                    "w": tile_w,
+                    "h": tile_h,
+                    "mask_box": None,
+                    "prompt": tile_prompt,
+                    "scale_factor": tile_scale,
+                }
+            )
+            tile_option = input("Do you want to define another tile? (y/N): ")
+            if tile_option == "" or tile_option.upper() == "N":
+                add_tile = False
+
+    return custom_settings
+
+
+def interactive_sdupscale(
+    config,
+    context,
+):
+    settings = config.lcm_diffusion_setting
+    settings.diffusion_task = DiffusionTask.image_to_image.value
+    settings.init_image = ""
+    source_path = ""
+    steps = settings.inference_steps
+
+    while True:
+        custom_upscale_settings = None
+        option = input("Edit custom SD Upscale settings? (y/N): ")
+        if option.upper() == "Y":
+            config.lcm_diffusion_setting.inference_steps = steps
+            custom_upscale_settings = interactive_sdupscale_settings(config)
+            if not custom_upscale_settings:
+                return
+            source_path = custom_upscale_settings["source_file"]
+        else:
+            new_path = input(f"Image path ({source_path}): ")
+            if new_path != "":
+                source_path = new_path
+            if source_path == "":
+                print("Error : You need to provide a file in SD Upscale mode")
+                return
+            settings.strength = user_value(
+                float,
+                f"SD Upscale strength ({settings.strength}): ",
+                settings.strength,
+            )
+            settings.inference_steps = int(steps / settings.strength + 1)
+
+        output_path = FastStableDiffusionPaths.get_upscale_filepath(
+            source_path,
+            2,
+            config.generated_images.format,
+        )
+        generate_upscaled_image(
+            config,
+            source_path,
+            settings.strength,
+            upscale_settings=custom_upscale_settings,
+            context=context,
+            tile_overlap=32 if settings.use_openvino else 16,
+            output_path=output_path,
+            image_format=config.generated_images.format,
+        )
+        user_input = input("Continue in SD Upscale mode? (Y/n): ")
+        if user_input.upper() == "N":
+            settings.inference_steps = steps
+            return

src/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 ""