Merge branch 'main' into monorimet-patch-2

.flake8

@@ -2,4 +2,4 @@
 count = 1
 show-source = 1
 select = E9,F63,F7,F82
-exclude = lit.cfg.py, apps/language_models/scripts/vicuna.py, apps/language_models/src/pipelines/minigpt4_pipeline.py, apps/language_models/langchain/h2oai_pipeline.py
+exclude = lit.cfg.py

.github/workflows/nightly.yml

@@ -51,11 +51,11 @@ jobs:
       run: |
         ./setup_venv.ps1
         $env:SHARK_PACKAGE_VERSION=${{ env.package_version }}
-        pip wheel -v -w dist . --pre -f https://download.pytorch.org/whl/nightly/cpu -f https://llvm.github.io/torch-mlir/package-index/ -f https://nod-ai.github.io/SRT/pip-release-links.html
+        pip wheel -v -w dist . --pre -f https://download.pytorch.org/whl/nightly/cpu -f https://llvm.github.io/torch-mlir/package-index/ -f https://nod-ai.github.io/SHARK-Runtime/pip-release-links.html
         python process_skipfiles.py
         pyinstaller .\apps\stable_diffusion\shark_sd.spec
-        mv ./dist/nodai_shark_studio.exe ./dist/nodai_shark_studio_${{ env.package_version_ }}.exe
-        signtool sign /f c:\g\shark_02152023.cer /fd certHash /csp "eToken Base Cryptographic Provider" /k "${{ secrets.CI_CERT }}" ./dist/nodai_shark_studio_${{ env.package_version_ }}.exe
+        mv ./dist/shark_sd.exe ./dist/nodai_shark_sd_${{ env.package_version_ }}.exe
+        signtool sign /f c:\g\shark_02152023.cer /csp "eToken Base Cryptographic Provider" /k "${{ secrets.CI_CERT }}" ./dist/nodai_shark_sd_${{ env.package_version_ }}.exe
   
     - name: Upload Release Assets
       id: upload-release-assets
@@ -104,7 +104,7 @@ jobs:
         echo "DATE=$(date +'%Y-%m-%d')" >> $GITHUB_ENV
         python -m pip install --upgrade pip
         python -m pip install flake8 pytest toml
-        if [ -f requirements.txt ]; then pip install -r requirements.txt -f https://llvm.github.io/torch-mlir/package-index/ -f https://nod-ai.github.io/SRT/pip-release-links.html; fi
+        if [ -f requirements.txt ]; then pip install -r requirements.txt -f https://llvm.github.io/torch-mlir/package-index/ -f https://nod-ai.github.io/SHARK-Runtime/pip-release-links.html; fi
     - name: Lint with flake8
       run: |
         # stop the build if there are Python syntax errors or undefined names
@@ -144,7 +144,7 @@ jobs:
         source shark.venv/bin/activate
         package_version="$(printf '%(%Y%m%d)T.${{ github.run_number }}')"
         SHARK_PACKAGE_VERSION=${package_version} \
-        pip wheel -v -w wheelhouse . --pre -f https://download.pytorch.org/whl/nightly/torch -f https://llvm.github.io/torch-mlir/package-index/ -f https://nod-ai.github.io/SRT/pip-release-links.html
+        pip wheel -v -w wheelhouse . --pre -f https://download.pytorch.org/whl/nightly/torch -f https://llvm.github.io/torch-mlir/package-index/ -f https://nod-ai.github.io/SHARK-Runtime/pip-release-links.html
         # Install the built wheel
         pip install ./wheelhouse/nodai*
         # Validate the Models

.github/workflows/test-models.yml

@@ -35,8 +35,6 @@ jobs:
         include:
           - os: ubuntu-latest
             suite: lint
-          - os: MacStudio
-            suite: metal
         exclude:
           - os: ubuntu-latest
             suite: vulkan
@@ -48,8 +46,6 @@ jobs:
             suite: cuda
           - os: MacStudio
             suite: cpu
-          - os: MacStudio
-            suite: vulkan
           - os: icelake
             suite: vulkan
           - os: icelake
@@ -112,42 +108,41 @@ jobs:
         cd $GITHUB_WORKSPACE
         PYTHON=python${{ matrix.python-version }} IMPORTER=1 ./setup_venv.sh
         source shark.venv/bin/activate
-        pytest --benchmark=native --update_tank -k cpu 
+        pytest --forked --benchmark=native --ci --ci_sha=${SHORT_SHA} --update_tank --tank_url="gs://shark_tank/nightly/" -k cpu 
         gsutil cp ./bench_results.csv gs://shark-public/builder/bench_results/${DATE}/bench_results_cpu_${SHORT_SHA}.csv
         gsutil cp gs://shark-public/builder/bench_results/${DATE}/bench_results_cpu_${SHORT_SHA}.csv gs://shark-public/builder/bench_results/latest/bench_results_cpu_latest.csv
-        python build_tools/vicuna_testing.py
 
     - name: Validate Models on NVIDIA GPU
       if: matrix.suite == 'cuda'
       run: |
         cd $GITHUB_WORKSPACE
-        PYTHON=python${{ matrix.python-version }} IMPORTER=1 ./setup_venv.sh
+        PYTHON=python${{ matrix.python-version }} ./setup_venv.sh
         source shark.venv/bin/activate
-        pytest --benchmark=native --update_tank -k cuda
+        pytest --forked --benchmark=native --ci --ci_sha=${SHORT_SHA} --update_tank --tank_url="gs://shark_tank/nightly/" -k cuda
         gsutil cp ./bench_results.csv gs://shark-public/builder/bench_results/${DATE}/bench_results_cuda_${SHORT_SHA}.csv
         gsutil cp gs://shark-public/builder/bench_results/${DATE}/bench_results_cuda_${SHORT_SHA}.csv gs://shark-public/builder/bench_results/latest/bench_results_cuda_latest.csv
         # Disabled due to black image bug
         # python build_tools/stable_diffusion_testing.py --device=cuda 
 
     - name: Validate Vulkan Models (MacOS)
-      if: matrix.suite == 'metal' && matrix.os == 'MacStudio'
+      if: matrix.suite == 'vulkan' && matrix.os == 'MacStudio'
       run: |
         cd $GITHUB_WORKSPACE
         PYTHON=python${{ matrix.python-version }} ./setup_venv.sh
         source shark.venv/bin/activate
+        export DYLD_LIBRARY_PATH=/usr/local/lib/
         echo $PATH
         pip list | grep -E "torch|iree"
-      # disabled due to a low-visibility memory issue with pytest on macos.
-      # pytest --ci --ci_sha=${SHORT_SHA} --local_tank_cache="/Volumes/builder/anush/shark_cache" --tank_url="gs://shark_tank/nightly/" -k metal
+        pytest --ci --ci_sha=${SHORT_SHA} --local_tank_cache="/Volumes/builder/anush/shark_cache" --tank_url="gs://shark_tank/nightly/" -k vulkan
 
     - name: Validate Vulkan Models (a100)
       if: matrix.suite == 'vulkan' && matrix.os == 'a100'
       run: |
         cd $GITHUB_WORKSPACE
-        PYTHON=python${{ matrix.python-version }} IMPORTER=1 ./setup_venv.sh
+        PYTHON=python${{ matrix.python-version }} ./setup_venv.sh
         source shark.venv/bin/activate
-        pytest --update_tank -k vulkan
-        python build_tools/stable_diffusion_testing.py --device=vulkan --no-exit_on_fail
+        pytest --forked --benchmark="native" --ci --ci_sha=${SHORT_SHA} --update_tank --tank_url="gs://shark_tank/nightly/" -k vulkan
+        python build_tools/stable_diffusion_testing.py --device=vulkan
 
     - name: Validate Vulkan Models (Windows)
       if: matrix.suite == 'vulkan' && matrix.os == '7950x'

.gitignore

@@ -2,8 +2,6 @@
 __pycache__/
 *.py[cod]
 *$py.class
-*.mlir
-*.vmfb
 
 # C extensions
 *.so
@@ -159,7 +157,7 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-.idea/
+#.idea/
 
 # vscode related
 .vscode
@@ -182,23 +180,10 @@ generated_imgs/
 
 # Custom model related artefacts
 variants.json
-/models/
+models/
 
 # models folder
 apps/stable_diffusion/web/models/
 
 # Stencil annotators.
 stencil_annotator/
-
-# For DocuChat
-apps/language_models/langchain/user_path/
-db_dir_UserData
-
-# Embeded browser cache and other
-apps/stable_diffusion/web/EBWebView/
-
-# Llama2 tokenizer configs
-llama2_tokenizer_configs/
-
-# Webview2 runtime artefacts
-EBWebView/

.gitmodules

@@ -1,4 +1,4 @@
 [submodule "inference/thirdparty/shark-runtime"]
 	path = inference/thirdparty/shark-runtime
-	url =https://github.com/nod-ai/SRT.git
+	url =https://github.com/nod-ai/SHARK-Runtime.git
 	branch = shark-06032022

README.md

@@ -10,7 +10,7 @@ High Performance Machine Learning Distribution
   <summary>Prerequisites - Drivers </summary>
   
 #### Install your Windows hardware drivers
-* [AMD RDNA Users] Download the latest driver (23.2.1 is the oldest supported) [here](https://www.amd.com/en/support).
+* [AMD RDNA Users] Download the latest driver [here](https://www.amd.com/en/support/kb/release-notes/rn-rad-win-23-2-1).
 * [macOS Users] Download and install the 1.3.216 Vulkan SDK from [here](https://sdk.lunarg.com/sdk/download/1.3.216.0/mac/vulkansdk-macos-1.3.216.0.dmg). Newer versions of the SDK will not work. 
 * [Nvidia Users] Download and install the latest CUDA / Vulkan drivers from [here](https://developer.nvidia.com/cuda-downloads)
   
@@ -170,7 +170,7 @@ python -m pip install --upgrade pip
 This step pip installs SHARK and related packages on Linux Python 3.8, 3.10 and 3.11 and macOS / Windows Python 3.11
 
 ```shell
-pip install nodai-shark -f https://nod-ai.github.io/SHARK/package-index/ -f https://llvm.github.io/torch-mlir/package-index/ -f  https://nod-ai.github.io/SRT/pip-release-links.html --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+pip install nodai-shark -f https://nod-ai.github.io/SHARK/package-index/ -f https://llvm.github.io/torch-mlir/package-index/ -f  https://nod-ai.github.io/SHARK-Runtime/pip-release-links.html --extra-index-url https://download.pytorch.org/whl/nightly/cpu
 ```
 
 ### Run shark tank model tests.
@@ -254,6 +254,7 @@ if you want to instead incorporate this into a python script, you can pass the `
 ```
 shark_module = SharkInference(
         mlir_model,
+        func_name,
         device=args.device,
         mlir_dialect="tm_tensor",
         dispatch_benchmarks="all",
@@ -296,7 +297,7 @@ torch_mlir, func_name = mlir_importer.import_mlir(tracing_required=True)
 # SharkInference accepts mlir in linalg, mhlo, and tosa dialect.
 
 from shark.shark_inference import SharkInference
-shark_module = SharkInference(torch_mlir, device="cpu", mlir_dialect="linalg")
+shark_module = SharkInference(torch_mlir, func_name, device="cpu", mlir_dialect="linalg")
 shark_module.compile()
 result = shark_module.forward((input))
 
@@ -319,17 +320,12 @@ mhlo_ir = r"""builtin.module  {
 
 arg0 = np.ones((1, 4)).astype(np.float32)
 arg1 = np.ones((4, 1)).astype(np.float32)
-shark_module = SharkInference(mhlo_ir, device="cpu", mlir_dialect="mhlo")
+shark_module = SharkInference(mhlo_ir, func_name="forward", device="cpu", mlir_dialect="mhlo")
 shark_module.compile()
 result = shark_module.forward((arg0, arg1))
 ```
 </details>
 
-## Examples Using the REST API
-
-* [Setting up SHARK for use with Blender](./docs/shark_sd_blender.md)
-* [Setting up SHARK for use with Koboldcpp](./docs/shark_sd_koboldcpp.md)
-
 ## Supported and Validated Models
 
 SHARK is maintained to support the latest innovations in ML Models: 

apps/language_models/README.md

@@ -1,16 +0,0 @@
-## CodeGen Setup using SHARK-server
-
-### Setup Server
-- clone SHARK and setup the venv
-- host the server using `python apps/stable_diffusion/web/index.py --api --server_port=<PORT>`
-- default server address is `http://0.0.0.0:8080`
-
-### Setup Client
-1. fauxpilot-vscode (VSCode Extension):
-- Code for the extension can be found [here](https://github.com/Venthe/vscode-fauxpilot)
-- PreReq: VSCode extension (will need [`nodejs` and `npm`](https://nodejs.org/en/download) to compile and run the extension)
-- Compile and Run the extension on VSCode (press F5 on VSCode), this opens a new VSCode window with the extension running
-- Open VSCode settings, search for fauxpilot in settings and modify `server : http://<IP>:<PORT>`, `Model : codegen` , `Max Lines : 30`
-
-2. Others (REST API curl, OpenAI Python bindings) as shown [here](https://github.com/fauxpilot/fauxpilot/blob/main/documentation/client.md)
-- using Github Copilot VSCode extension with SHARK-server needs more work to be functional.

apps/language_models/langchain/README.md

@@ -1,18 +0,0 @@
-# Langchain
-
-## How to run the model
-
-1.) Install all the dependencies by running:
-```shell
-pip install -r apps/language_models/langchain/langchain_requirements.txt
-sudo apt-get install -y libmagic-dev poppler-utils tesseract-ocr libtesseract-dev libreoffice
-```
-
-2.) Create a folder named `user_path` in `apps/language_models/langchain/` directory.
-
-Now, you are ready to use the model.
-
-3.) To run the model, run the following command:
-```shell
-python apps/language_models/langchain/gen.py --cli=True
-```

apps/language_models/langchain/cli.py

@@ -1,186 +0,0 @@
-import copy
-import torch
-
-from evaluate_params import eval_func_param_names
-from gen import Langchain
-from prompter import non_hf_types
-from utils import clear_torch_cache, NullContext, get_kwargs
-
-
-def run_cli(  # for local function:
-    base_model=None,
-    lora_weights=None,
-    inference_server=None,
-    debug=None,
-    chat_context=None,
-    examples=None,
-    memory_restriction_level=None,
-    # for get_model:
-    score_model=None,
-    load_8bit=None,
-    load_4bit=None,
-    load_half=None,
-    load_gptq=None,
-    use_safetensors=None,
-    infer_devices=None,
-    tokenizer_base_model=None,
-    gpu_id=None,
-    local_files_only=None,
-    resume_download=None,
-    use_auth_token=None,
-    trust_remote_code=None,
-    offload_folder=None,
-    compile_model=None,
-    # for some evaluate args
-    stream_output=None,
-    prompt_type=None,
-    prompt_dict=None,
-    temperature=None,
-    top_p=None,
-    top_k=None,
-    num_beams=None,
-    max_new_tokens=None,
-    min_new_tokens=None,
-    early_stopping=None,
-    max_time=None,
-    repetition_penalty=None,
-    num_return_sequences=None,
-    do_sample=None,
-    chat=None,
-    langchain_mode=None,
-    langchain_action=None,
-    document_choice=None,
-    top_k_docs=None,
-    chunk=None,
-    chunk_size=None,
-    # for evaluate kwargs
-    src_lang=None,
-    tgt_lang=None,
-    concurrency_count=None,
-    save_dir=None,
-    sanitize_bot_response=None,
-    model_state0=None,
-    max_max_new_tokens=None,
-    is_public=None,
-    max_max_time=None,
-    raise_generate_gpu_exceptions=None,
-    load_db_if_exists=None,
-    dbs=None,
-    user_path=None,
-    detect_user_path_changes_every_query=None,
-    use_openai_embedding=None,
-    use_openai_model=None,
-    hf_embedding_model=None,
-    db_type=None,
-    n_jobs=None,
-    first_para=None,
-    text_limit=None,
-    verbose=None,
-    cli=None,
-    reverse_docs=None,
-    use_cache=None,
-    auto_reduce_chunks=None,
-    max_chunks=None,
-    model_lock=None,
-    force_langchain_evaluate=None,
-    model_state_none=None,
-    # unique to this function:
-    cli_loop=None,
-):
-    Langchain.check_locals(**locals())
-
-    score_model = ""  # FIXME: For now, so user doesn't have to pass
-    n_gpus = torch.cuda.device_count() if torch.cuda.is_available else 0
-    device = "cpu" if n_gpus == 0 else "cuda"
-    context_class = NullContext if n_gpus > 1 or n_gpus == 0 else torch.device
-
-    with context_class(device):
-        from functools import partial
-
-        # get score model
-        smodel, stokenizer, sdevice = Langchain.get_score_model(
-            reward_type=True,
-            **get_kwargs(
-                Langchain.get_score_model,
-                exclude_names=["reward_type"],
-                **locals()
-            )
-        )
-
-        model, tokenizer, device = Langchain.get_model(
-            reward_type=False,
-            **get_kwargs(
-                Langchain.get_model, exclude_names=["reward_type"], **locals()
-            )
-        )
-        model_dict = dict(
-            base_model=base_model,
-            tokenizer_base_model=tokenizer_base_model,
-            lora_weights=lora_weights,
-            inference_server=inference_server,
-            prompt_type=prompt_type,
-            prompt_dict=prompt_dict,
-        )
-        model_state = dict(model=model, tokenizer=tokenizer, device=device)
-        model_state.update(model_dict)
-        my_db_state = [None]
-        fun = partial(
-            Langchain.evaluate,
-            model_state,
-            my_db_state,
-            **get_kwargs(
-                Langchain.evaluate,
-                exclude_names=["model_state", "my_db_state"]
-                + eval_func_param_names,
-                **locals()
-            )
-        )
-
-        example1 = examples[-1]  # pick reference example
-        all_generations = []
-        while True:
-            clear_torch_cache()
-            instruction = input("\nEnter an instruction: ")
-            if instruction == "exit":
-                break
-
-            eval_vars = copy.deepcopy(example1)
-            eval_vars[eval_func_param_names.index("instruction")] = eval_vars[
-                eval_func_param_names.index("instruction_nochat")
-            ] = instruction
-            eval_vars[eval_func_param_names.index("iinput")] = eval_vars[
-                eval_func_param_names.index("iinput_nochat")
-            ] = ""  # no input yet
-            eval_vars[
-                eval_func_param_names.index("context")
-            ] = ""  # no context yet
-
-            # grab other parameters, like langchain_mode
-            for k in eval_func_param_names:
-                if k in locals():
-                    eval_vars[eval_func_param_names.index(k)] = locals()[k]
-
-            gener = fun(*tuple(eval_vars))
-            outr = ""
-            res_old = ""
-            for gen_output in gener:
-                res = gen_output["response"]
-                extra = gen_output["sources"]
-                if base_model not in non_hf_types or base_model in ["llama"]:
-                    if not stream_output:
-                        print(res)
-                    else:
-                        # then stream output for gradio that has full output each generation, so need here to show only new chars
-                        diff = res[len(res_old) :]
-                        print(diff, end="", flush=True)
-                        res_old = res
-                    outr = res  # don't accumulate
-                else:
-                    outr += res  # just is one thing
-                    if extra:
-                        # show sources at end after model itself had streamed to std rest of response
-                        print(extra, flush=True)
-            all_generations.append(outr + "\n")
-            if not cli_loop:
-                break
-    return all_generations

apps/language_models/langchain/enums.py

@@ -1,103 +0,0 @@
-from enum import Enum
-
-
-class PromptType(Enum):
-    custom = -1
-    plain = 0
-    instruct = 1
-    quality = 2
-    human_bot = 3
-    dai_faq = 4
-    summarize = 5
-    simple_instruct = 6
-    instruct_vicuna = 7
-    instruct_with_end = 8
-    human_bot_orig = 9
-    prompt_answer = 10
-    open_assistant = 11
-    wizard_lm = 12
-    wizard_mega = 13
-    instruct_vicuna2 = 14
-    instruct_vicuna3 = 15
-    wizard2 = 16
-    wizard3 = 17
-    instruct_simple = 18
-    wizard_vicuna = 19
-    openai = 20
-    openai_chat = 21
-    gptj = 22
-    prompt_answer_openllama = 23
-    vicuna11 = 24
-    mptinstruct = 25
-    mptchat = 26
-    falcon = 27
-
-
-class DocumentChoices(Enum):
-    All_Relevant = 0
-    All_Relevant_Only_Sources = 1
-    Only_All_Sources = 2
-    Just_LLM = 3
-
-
-non_query_commands = [
-    DocumentChoices.All_Relevant_Only_Sources.name,
-    DocumentChoices.Only_All_Sources.name,
-]
-
-
-class LangChainMode(Enum):
-    """LangChain mode"""
-
-    DISABLED = "Disabled"
-    CHAT_LLM = "ChatLLM"
-    LLM = "LLM"
-    ALL = "All"
-    WIKI = "wiki"
-    WIKI_FULL = "wiki_full"
-    USER_DATA = "UserData"
-    MY_DATA = "MyData"
-    GITHUB_H2OGPT = "github h2oGPT"
-    H2O_DAI_DOCS = "DriverlessAI docs"
-
-
-class LangChainAction(Enum):
-    """LangChain action"""
-
-    QUERY = "Query"
-    # WIP:
-    # SUMMARIZE_MAP = "Summarize_map_reduce"
-    SUMMARIZE_MAP = "Summarize"
-    SUMMARIZE_ALL = "Summarize_all"
-    SUMMARIZE_REFINE = "Summarize_refine"
-
-
-no_server_str = no_lora_str = no_model_str = "[None/Remove]"
-
-# from site-packages/langchain/llms/openai.py
-# but needed since ChatOpenAI doesn't have this information
-model_token_mapping = {
-    "gpt-4": 8192,
-    "gpt-4-0314": 8192,
-    "gpt-4-32k": 32768,
-    "gpt-4-32k-0314": 32768,
-    "gpt-3.5-turbo": 4096,
-    "gpt-3.5-turbo-16k": 16 * 1024,
-    "gpt-3.5-turbo-0301": 4096,
-    "text-ada-001": 2049,
-    "ada": 2049,
-    "text-babbage-001": 2040,
-    "babbage": 2049,
-    "text-curie-001": 2049,
-    "curie": 2049,
-    "davinci": 2049,
-    "text-davinci-003": 4097,
-    "text-davinci-002": 4097,
-    "code-davinci-002": 8001,
-    "code-davinci-001": 8001,
-    "code-cushman-002": 2048,
-    "code-cushman-001": 2048,
-}
-
-source_prefix = "Sources [Score | Link]:"
-source_postfix = "End Sources<p>"

apps/language_models/langchain/evaluate_params.py

@@ -1,53 +0,0 @@
-no_default_param_names = [
-    "instruction",
-    "iinput",
-    "context",
-    "instruction_nochat",
-    "iinput_nochat",
-]
-
-gen_hyper = [
-    "temperature",
-    "top_p",
-    "top_k",
-    "num_beams",
-    "max_new_tokens",
-    "min_new_tokens",
-    "early_stopping",
-    "max_time",
-    "repetition_penalty",
-    "num_return_sequences",
-    "do_sample",
-]
-
-eval_func_param_names = (
-    [
-        "instruction",
-        "iinput",
-        "context",
-        "stream_output",
-        "prompt_type",
-        "prompt_dict",
-    ]
-    + gen_hyper
-    + [
-        "chat",
-        "instruction_nochat",
-        "iinput_nochat",
-        "langchain_mode",
-        "langchain_action",
-        "top_k_docs",
-        "chunk",
-        "chunk_size",
-        "document_choice",
-    ]
-)
-
-# form evaluate defaults for submit_nochat_api
-eval_func_param_names_defaults = eval_func_param_names.copy()
-for k in no_default_param_names:
-    if k in eval_func_param_names_defaults:
-        eval_func_param_names_defaults.remove(k)
-
-
-eval_extra_columns = ["prompt", "response", "score"]

apps/language_models/langchain/expanded_pipelines.py

@@ -1,846 +0,0 @@
-from __future__ import annotations
-from typing import (
-    Any,
-    Mapping,
-    Optional,
-    Dict,
-    List,
-    Sequence,
-    Tuple,
-    Union,
-    Protocol,
-)
-import inspect
-import json
-import warnings
-from pathlib import Path
-import yaml
-from abc import ABC, abstractmethod
-import langchain
-from langchain.base_language import BaseLanguageModel
-from langchain.callbacks.base import BaseCallbackManager
-from langchain.chains.question_answering import stuff_prompt
-from langchain.prompts.base import BasePromptTemplate
-from langchain.docstore.document import Document
-from langchain.callbacks.manager import (
-    CallbackManager,
-    CallbackManagerForChainRun,
-    Callbacks,
-)
-from langchain.load.serializable import Serializable
-from langchain.schema import RUN_KEY, BaseMemory, RunInfo
-from langchain.input import get_colored_text
-from langchain.load.dump import dumpd
-from langchain.prompts.prompt import PromptTemplate
-from langchain.schema import LLMResult, PromptValue
-from pydantic import Extra, Field, root_validator, validator
-
-
-def _get_verbosity() -> bool:
-    return langchain.verbose
-
-
-def format_document(doc: Document, prompt: BasePromptTemplate) -> str:
-    """Format a document into a string based on a prompt template."""
-    base_info = {"page_content": doc.page_content}
-    base_info.update(doc.metadata)
-    missing_metadata = set(prompt.input_variables).difference(base_info)
-    if len(missing_metadata) > 0:
-        required_metadata = [
-            iv for iv in prompt.input_variables if iv != "page_content"
-        ]
-        raise ValueError(
-            f"Document prompt requires documents to have metadata variables: "
-            f"{required_metadata}. Received document with missing metadata: "
-            f"{list(missing_metadata)}."
-        )
-    document_info = {k: base_info[k] for k in prompt.input_variables}
-    return prompt.format(**document_info)
-
-
-class Chain(Serializable, ABC):
-    """Base interface that all chains should implement."""
-
-    memory: Optional[BaseMemory] = None
-    callbacks: Callbacks = Field(default=None, exclude=True)
-    callback_manager: Optional[BaseCallbackManager] = Field(
-        default=None, exclude=True
-    )
-    verbose: bool = Field(
-        default_factory=_get_verbosity
-    )  # Whether to print the response text
-    tags: Optional[List[str]] = None
-
-    class Config:
-        """Configuration for this pydantic object."""
-
-        arbitrary_types_allowed = True
-
-    @property
-    def _chain_type(self) -> str:
-        raise NotImplementedError("Saving not supported for this chain type.")
-
-    @root_validator()
-    def raise_deprecation(cls, values: Dict) -> Dict:
-        """Raise deprecation warning if callback_manager is used."""
-        if values.get("callback_manager") is not None:
-            warnings.warn(
-                "callback_manager is deprecated. Please use callbacks instead.",
-                DeprecationWarning,
-            )
-            values["callbacks"] = values.pop("callback_manager", None)
-        return values
-
-    @validator("verbose", pre=True, always=True)
-    def set_verbose(cls, verbose: Optional[bool]) -> bool:
-        """If verbose is None, set it.
-
-        This allows users to pass in None as verbose to access the global setting.
-        """
-        if verbose is None:
-            return _get_verbosity()
-        else:
-            return verbose
-
-    @property
-    @abstractmethod
-    def input_keys(self) -> List[str]:
-        """Input keys this chain expects."""
-
-    @property
-    @abstractmethod
-    def output_keys(self) -> List[str]:
-        """Output keys this chain expects."""
-
-    def _validate_inputs(self, inputs: Dict[str, Any]) -> None:
-        """Check that all inputs are present."""
-        missing_keys = set(self.input_keys).difference(inputs)
-        if missing_keys:
-            raise ValueError(f"Missing some input keys: {missing_keys}")
-
-    def _validate_outputs(self, outputs: Dict[str, Any]) -> None:
-        missing_keys = set(self.output_keys).difference(outputs)
-        if missing_keys:
-            raise ValueError(f"Missing some output keys: {missing_keys}")
-
-    @abstractmethod
-    def _call(
-        self,
-        inputs: Dict[str, Any],
-        run_manager: Optional[CallbackManagerForChainRun] = None,
-    ) -> Dict[str, Any]:
-        """Run the logic of this chain and return the output."""
-
-    def __call__(
-        self,
-        inputs: Union[Dict[str, Any], Any],
-        return_only_outputs: bool = False,
-        callbacks: Callbacks = None,
-        *,
-        tags: Optional[List[str]] = None,
-        include_run_info: bool = False,
-    ) -> Dict[str, Any]:
-        """Run the logic of this chain and add to output if desired.
-
-        Args:
-            inputs: Dictionary of inputs, or single input if chain expects
-                only one param.
-            return_only_outputs: boolean for whether to return only outputs in the
-                response. If True, only new keys generated by this chain will be
-                returned. If False, both input keys and new keys generated by this
-                chain will be returned. Defaults to False.
-            callbacks: Callbacks to use for this chain run. If not provided, will
-                use the callbacks provided to the chain.
-            include_run_info: Whether to include run info in the response. Defaults
-                to False.
-        """
-        input_docs = inputs["input_documents"]
-        missing_keys = set(self.input_keys).difference(inputs)
-        if missing_keys:
-            raise ValueError(f"Missing some input keys: {missing_keys}")
-
-        callback_manager = CallbackManager.configure(
-            callbacks, self.callbacks, self.verbose, tags, self.tags
-        )
-        run_manager = callback_manager.on_chain_start(
-            dumpd(self),
-            inputs,
-        )
-
-        if "is_first" in inputs.keys() and not inputs["is_first"]:
-            run_manager_ = run_manager
-            input_list = [inputs]
-            stop = None
-            prompts = []
-            for inputs in input_list:
-                selected_inputs = {
-                    k: inputs[k] for k in self.prompt.input_variables
-                }
-                prompt = self.prompt.format_prompt(**selected_inputs)
-                _colored_text = get_colored_text(prompt.to_string(), "green")
-                _text = "Prompt after formatting:\n" + _colored_text
-                if run_manager_:
-                    run_manager_.on_text(_text, end="\n", verbose=self.verbose)
-                if "stop" in inputs and inputs["stop"] != stop:
-                    raise ValueError(
-                        "If `stop` is present in any inputs, should be present in all."
-                    )
-                prompts.append(prompt)
-
-            prompt_strings = [p.to_string() for p in prompts]
-            prompts = prompt_strings
-            callbacks = run_manager_.get_child() if run_manager_ else None
-            tags = None
-
-            """Run the LLM on the given prompt and input."""
-            # If string is passed in directly no errors will be raised but outputs will
-            # not make sense.
-            if not isinstance(prompts, list):
-                raise ValueError(
-                    "Argument 'prompts' is expected to be of type List[str], received"
-                    f" argument of type {type(prompts)}."
-                )
-            params = self.llm.dict()
-            params["stop"] = stop
-            options = {"stop": stop}
-            disregard_cache = self.llm.cache is not None and not self.llm.cache
-            callback_manager = CallbackManager.configure(
-                callbacks,
-                self.llm.callbacks,
-                self.llm.verbose,
-                tags,
-                self.llm.tags,
-            )
-            if langchain.llm_cache is None or disregard_cache:
-                # This happens when langchain.cache is None, but self.cache is True
-                if self.llm.cache is not None and self.cache:
-                    raise ValueError(
-                        "Asked to cache, but no cache found at `langchain.cache`."
-                    )
-                run_manager_ = callback_manager.on_llm_start(
-                    dumpd(self),
-                    prompts,
-                    invocation_params=params,
-                    options=options,
-                )
-
-                generations = []
-                for prompt in prompts:
-                    inputs_ = prompt
-                    num_workers = None
-                    batch_size = None
-
-                    if num_workers is None:
-                        if self.llm.pipeline._num_workers is None:
-                            num_workers = 0
-                        else:
-                            num_workers = self.llm.pipeline._num_workers
-                    if batch_size is None:
-                        if self.llm.pipeline._batch_size is None:
-                            batch_size = 1
-                        else:
-                            batch_size = self.llm.pipeline._batch_size
-
-                    preprocess_params = {}
-                    generate_kwargs = {}
-                    preprocess_params.update(generate_kwargs)
-                    forward_params = generate_kwargs
-                    postprocess_params = {}
-                    # Fuse __init__ params and __call__ params without modifying the __init__ ones.
-                    preprocess_params = {
-                        **self.llm.pipeline._preprocess_params,
-                        **preprocess_params,
-                    }
-                    forward_params = {
-                        **self.llm.pipeline._forward_params,
-                        **forward_params,
-                    }
-                    postprocess_params = {
-                        **self.llm.pipeline._postprocess_params,
-                        **postprocess_params,
-                    }
-
-                    self.llm.pipeline.call_count += 1
-                    if (
-                        self.llm.pipeline.call_count > 10
-                        and self.llm.pipeline.framework == "pt"
-                        and self.llm.pipeline.device.type == "cuda"
-                    ):
-                        warnings.warn(
-                            "You seem to be using the pipelines sequentially on GPU. In order to maximize efficiency please use a"
-                            " dataset",
-                            UserWarning,
-                        )
-
-                    model_inputs = self.llm.pipeline.preprocess(
-                        inputs_, **preprocess_params
-                    )
-                    model_outputs = self.llm.pipeline.forward(
-                        model_inputs, **forward_params
-                    )
-                    model_outputs["process"] = False
-                    return model_outputs
-                output = LLMResult(generations=generations)
-                run_manager_.on_llm_end(output)
-                if run_manager_:
-                    output.run = RunInfo(run_id=run_manager_.run_id)
-                response = output
-
-            outputs = [
-                # Get the text of the top generated string.
-                {self.output_key: generation[0].text}
-                for generation in response.generations
-            ][0]
-            run_manager.on_chain_end(outputs)
-            final_outputs: Dict[str, Any] = self.prep_outputs(
-                inputs, outputs, return_only_outputs
-            )
-            if include_run_info:
-                final_outputs[RUN_KEY] = RunInfo(run_id=run_manager.run_id)
-            return final_outputs
-        else:
-            _run_manager = (
-                run_manager or CallbackManagerForChainRun.get_noop_manager()
-            )
-            docs = inputs[self.input_key]
-            # Other keys are assumed to be needed for LLM prediction
-            other_keys = {
-                k: v for k, v in inputs.items() if k != self.input_key
-            }
-            doc_strings = [
-                format_document(doc, self.document_prompt) for doc in docs
-            ]
-            # Join the documents together to put them in the prompt.
-            inputs = {
-                k: v
-                for k, v in other_keys.items()
-                if k in self.llm_chain.prompt.input_variables
-            }
-            inputs[self.document_variable_name] = self.document_separator.join(
-                doc_strings
-            )
-            inputs["is_first"] = False
-            inputs["input_documents"] = input_docs
-
-            # Call predict on the LLM.
-            output = self.llm_chain(inputs, callbacks=_run_manager.get_child())
-            if "process" in output.keys() and not output["process"]:
-                return output
-            output = output[self.llm_chain.output_key]
-            extra_return_dict = {}
-        extra_return_dict[self.output_key] = output
-        outputs = extra_return_dict
-        run_manager.on_chain_end(outputs)
-        final_outputs: Dict[str, Any] = self.prep_outputs(
-            inputs, outputs, return_only_outputs
-        )
-        if include_run_info:
-            final_outputs[RUN_KEY] = RunInfo(run_id=run_manager.run_id)
-        return final_outputs
-
-    def prep_outputs(
-        self,
-        inputs: Dict[str, str],
-        outputs: Dict[str, str],
-        return_only_outputs: bool = False,
-    ) -> Dict[str, str]:
-        """Validate and prep outputs."""
-        self._validate_outputs(outputs)
-        if self.memory is not None:
-            self.memory.save_context(inputs, outputs)
-        if return_only_outputs:
-            return outputs
-        else:
-            return {**inputs, **outputs}
-
-    def prep_inputs(
-        self, inputs: Union[Dict[str, Any], Any]
-    ) -> Dict[str, str]:
-        """Validate and prep inputs."""
-        if not isinstance(inputs, dict):
-            _input_keys = set(self.input_keys)
-            if self.memory is not None:
-                # If there are multiple input keys, but some get set by memory so that
-                # only one is not set, we can still figure out which key it is.
-                _input_keys = _input_keys.difference(
-                    self.memory.memory_variables
-                )
-            if len(_input_keys) != 1:
-                raise ValueError(
-                    f"A single string input was passed in, but this chain expects "
-                    f"multiple inputs ({_input_keys}). When a chain expects "
-                    f"multiple inputs, please call it by passing in a dictionary, "
-                    "eg `chain({'foo': 1, 'bar': 2})`"
-                )
-            inputs = {list(_input_keys)[0]: inputs}
-        if self.memory is not None:
-            external_context = self.memory.load_memory_variables(inputs)
-            inputs = dict(inputs, **external_context)
-        self._validate_inputs(inputs)
-        return inputs
-
-    def apply(
-        self, input_list: List[Dict[str, Any]], callbacks: Callbacks = None
-    ) -> List[Dict[str, str]]:
-        """Call the chain on all inputs in the list."""
-        return [self(inputs, callbacks=callbacks) for inputs in input_list]
-
-    def run(
-        self,
-        *args: Any,
-        callbacks: Callbacks = None,
-        tags: Optional[List[str]] = None,
-        **kwargs: Any,
-    ) -> str:
-        """Run the chain as text in, text out or multiple variables, text out."""
-        if len(self.output_keys) != 1:
-            raise ValueError(
-                f"`run` not supported when there is not exactly "
-                f"one output key. Got {self.output_keys}."
-            )
-
-        if args and not kwargs:
-            if len(args) != 1:
-                raise ValueError(
-                    "`run` supports only one positional argument."
-                )
-            return self(args[0], callbacks=callbacks, tags=tags)[
-                self.output_keys[0]
-            ]
-
-        if kwargs and not args:
-            return self(kwargs, callbacks=callbacks, tags=tags)[
-                self.output_keys[0]
-            ]
-
-        if not kwargs and not args:
-            raise ValueError(
-                "`run` supported with either positional arguments or keyword arguments,"
-                " but none were provided."
-            )
-
-        raise ValueError(
-            f"`run` supported with either positional arguments or keyword arguments"
-            f" but not both. Got args: {args} and kwargs: {kwargs}."
-        )
-
-    def dict(self, **kwargs: Any) -> Dict:
-        """Return dictionary representation of chain."""
-        if self.memory is not None:
-            raise ValueError("Saving of memory is not yet supported.")
-        _dict = super().dict()
-        _dict["_type"] = self._chain_type
-        return _dict
-
-    def save(self, file_path: Union[Path, str]) -> None:
-        """Save the chain.
-
-        Args:
-            file_path: Path to file to save the chain to.
-
-        Example:
-        .. code-block:: python
-
-            chain.save(file_path="path/chain.yaml")
-        """
-        # Convert file to Path object.
-        if isinstance(file_path, str):
-            save_path = Path(file_path)
-        else:
-            save_path = file_path
-
-        directory_path = save_path.parent
-        directory_path.mkdir(parents=True, exist_ok=True)
-
-        # Fetch dictionary to save
-        chain_dict = self.dict()
-
-        if save_path.suffix == ".json":
-            with open(file_path, "w") as f:
-                json.dump(chain_dict, f, indent=4)
-        elif save_path.suffix == ".yaml":
-            with open(file_path, "w") as f:
-                yaml.dump(chain_dict, f, default_flow_style=False)
-        else:
-            raise ValueError(f"{save_path} must be json or yaml")
-
-
-class BaseCombineDocumentsChain(Chain, ABC):
-    """Base interface for chains combining documents."""
-
-    input_key: str = "input_documents"  #: :meta private:
-    output_key: str = "output_text"  #: :meta private:
-
-    @property
-    def input_keys(self) -> List[str]:
-        """Expect input key.
-
-        :meta private:
-        """
-        return [self.input_key]
-
-    @property
-    def output_keys(self) -> List[str]:
-        """Return output key.
-
-        :meta private:
-        """
-        return [self.output_key]
-
-    def prompt_length(
-        self, docs: List[Document], **kwargs: Any
-    ) -> Optional[int]:
-        """Return the prompt length given the documents passed in.
-
-        Returns None if the method does not depend on the prompt length.
-        """
-        return None
-
-    def _call(
-        self,
-        inputs: Dict[str, List[Document]],
-        run_manager: Optional[CallbackManagerForChainRun] = None,
-    ) -> Dict[str, str]:
-        _run_manager = (
-            run_manager or CallbackManagerForChainRun.get_noop_manager()
-        )
-        docs = inputs[self.input_key]
-        # Other keys are assumed to be needed for LLM prediction
-        other_keys = {k: v for k, v in inputs.items() if k != self.input_key}
-        doc_strings = [
-            format_document(doc, self.document_prompt) for doc in docs
-        ]
-        # Join the documents together to put them in the prompt.
-        inputs = {
-            k: v
-            for k, v in other_keys.items()
-            if k in self.llm_chain.prompt.input_variables
-        }
-        inputs[self.document_variable_name] = self.document_separator.join(
-            doc_strings
-        )
-
-        # Call predict on the LLM.
-        output, extra_return_dict = (
-            self.llm_chain(inputs, callbacks=_run_manager.get_child())[
-                self.llm_chain.output_key
-            ],
-            {},
-        )
-
-        extra_return_dict[self.output_key] = output
-        return extra_return_dict
-
-
-from pydantic import BaseModel
-
-
-class Generation(Serializable):
-    """Output of a single generation."""
-
-    text: str
-    """Generated text output."""
-
-    generation_info: Optional[Dict[str, Any]] = None
-    """Raw generation info response from the provider"""
-    """May include things like reason for finishing (e.g. in OpenAI)"""
-    # TODO: add log probs
-
-
-VALID_TASKS = ("text2text-generation", "text-generation", "summarization")
-
-
-class LLMChain(Chain):
-    """Chain to run queries against LLMs.
-
-    Example:
-        .. code-block:: python
-
-            from langchain import LLMChain, OpenAI, PromptTemplate
-            prompt_template = "Tell me a {adjective} joke"
-            prompt = PromptTemplate(
-                input_variables=["adjective"], template=prompt_template
-            )
-            llm = LLMChain(llm=OpenAI(), prompt=prompt)
-    """
-
-    @property
-    def lc_serializable(self) -> bool:
-        return True
-
-    prompt: BasePromptTemplate
-    """Prompt object to use."""
-    llm: BaseLanguageModel
-    output_key: str = "text"  #: :meta private:
-
-    class Config:
-        """Configuration for this pydantic object."""
-
-        extra = Extra.forbid
-        arbitrary_types_allowed = True
-
-    @property
-    def input_keys(self) -> List[str]:
-        """Will be whatever keys the prompt expects.
-
-        :meta private:
-        """
-        return self.prompt.input_variables
-
-    @property
-    def output_keys(self) -> List[str]:
-        """Will always return text key.
-
-        :meta private:
-        """
-        return [self.output_key]
-
-    def _call(
-        self,
-        inputs: Dict[str, Any],
-        run_manager: Optional[CallbackManagerForChainRun] = None,
-    ) -> Dict[str, str]:
-        prompts, stop = self.prep_prompts([inputs], run_manager=run_manager)
-        response = self.llm.generate_prompt(
-            prompts,
-            stop,
-            callbacks=run_manager.get_child() if run_manager else None,
-        )
-        return self.create_outputs(response)[0]
-
-    def prep_prompts(
-        self,
-        input_list: List[Dict[str, Any]],
-        run_manager: Optional[CallbackManagerForChainRun] = None,
-    ) -> Tuple[List[PromptValue], Optional[List[str]]]:
-        """Prepare prompts from inputs."""
-        stop = None
-        if "stop" in input_list[0]:
-            stop = input_list[0]["stop"]
-        prompts = []
-        for inputs in input_list:
-            selected_inputs = {
-                k: inputs[k] for k in self.prompt.input_variables
-            }
-            prompt = self.prompt.format_prompt(**selected_inputs)
-            _colored_text = get_colored_text(prompt.to_string(), "green")
-            _text = "Prompt after formatting:\n" + _colored_text
-            if run_manager:
-                run_manager.on_text(_text, end="\n", verbose=self.verbose)
-            if "stop" in inputs and inputs["stop"] != stop:
-                raise ValueError(
-                    "If `stop` is present in any inputs, should be present in all."
-                )
-            prompts.append(prompt)
-        return prompts, stop
-
-    def apply(
-        self, input_list: List[Dict[str, Any]], callbacks: Callbacks = None
-    ) -> List[Dict[str, str]]:
-        """Utilize the LLM generate method for speed gains."""
-        callback_manager = CallbackManager.configure(
-            callbacks, self.callbacks, self.verbose
-        )
-        run_manager = callback_manager.on_chain_start(
-            dumpd(self),
-            {"input_list": input_list},
-        )
-        try:
-            response = self.generate(input_list, run_manager=run_manager)
-        except (KeyboardInterrupt, Exception) as e:
-            run_manager.on_chain_error(e)
-            raise e
-        outputs = self.create_outputs(response)
-        run_manager.on_chain_end({"outputs": outputs})
-        return outputs
-
-    def create_outputs(self, response: LLMResult) -> List[Dict[str, str]]:
-        """Create outputs from response."""
-        return [
-            # Get the text of the top generated string.
-            {self.output_key: generation[0].text}
-            for generation in response.generations
-        ]
-
-    def predict_and_parse(
-        self, callbacks: Callbacks = None, **kwargs: Any
-    ) -> Union[str, List[str], Dict[str, Any]]:
-        """Call predict and then parse the results."""
-        result = self.predict(callbacks=callbacks, **kwargs)
-        if self.prompt.output_parser is not None:
-            return self.prompt.output_parser.parse(result)
-        else:
-            return result
-
-    def apply_and_parse(
-        self, input_list: List[Dict[str, Any]], callbacks: Callbacks = None
-    ) -> Sequence[Union[str, List[str], Dict[str, str]]]:
-        """Call apply and then parse the results."""
-        result = self.apply(input_list, callbacks=callbacks)
-        return self._parse_result(result)
-
-    def _parse_result(
-        self, result: List[Dict[str, str]]
-    ) -> Sequence[Union[str, List[str], Dict[str, str]]]:
-        if self.prompt.output_parser is not None:
-            return [
-                self.prompt.output_parser.parse(res[self.output_key])
-                for res in result
-            ]
-        else:
-            return result
-
-    @property
-    def _chain_type(self) -> str:
-        return "llm_chain"
-
-    @classmethod
-    def from_string(cls, llm: BaseLanguageModel, template: str) -> LLMChain:
-        """Create LLMChain from LLM and template."""
-        prompt_template = PromptTemplate.from_template(template)
-        return cls(llm=llm, prompt=prompt_template)
-
-
-def _get_default_document_prompt() -> PromptTemplate:
-    return PromptTemplate(
-        input_variables=["page_content"], template="{page_content}"
-    )
-
-
-class StuffDocumentsChain(BaseCombineDocumentsChain):
-    """Chain that combines documents by stuffing into context."""
-
-    llm_chain: LLMChain
-    """LLM wrapper to use after formatting documents."""
-    document_prompt: BasePromptTemplate = Field(
-        default_factory=_get_default_document_prompt
-    )
-    """Prompt to use to format each document."""
-    document_variable_name: str
-    """The variable name in the llm_chain to put the documents in.
-    If only one variable in the llm_chain, this need not be provided."""
-    document_separator: str = "\n\n"
-    """The string with which to join the formatted documents"""
-
-    class Config:
-        """Configuration for this pydantic object."""
-
-        extra = Extra.forbid
-        arbitrary_types_allowed = True
-
-    @root_validator(pre=True)
-    def get_default_document_variable_name(cls, values: Dict) -> Dict:
-        """Get default document variable name, if not provided."""
-        llm_chain_variables = values["llm_chain"].prompt.input_variables
-        if "document_variable_name" not in values:
-            if len(llm_chain_variables) == 1:
-                values["document_variable_name"] = llm_chain_variables[0]
-            else:
-                raise ValueError(
-                    "document_variable_name must be provided if there are "
-                    "multiple llm_chain_variables"
-                )
-        else:
-            if values["document_variable_name"] not in llm_chain_variables:
-                raise ValueError(
-                    f"document_variable_name {values['document_variable_name']} was "
-                    f"not found in llm_chain input_variables: {llm_chain_variables}"
-                )
-        return values
-
-    def _get_inputs(self, docs: List[Document], **kwargs: Any) -> dict:
-        # Format each document according to the prompt
-        doc_strings = [
-            format_document(doc, self.document_prompt) for doc in docs
-        ]
-        # Join the documents together to put them in the prompt.
-        inputs = {
-            k: v
-            for k, v in kwargs.items()
-            if k in self.llm_chain.prompt.input_variables
-        }
-        inputs[self.document_variable_name] = self.document_separator.join(
-            doc_strings
-        )
-        return inputs
-
-    def prompt_length(
-        self, docs: List[Document], **kwargs: Any
-    ) -> Optional[int]:
-        """Get the prompt length by formatting the prompt."""
-        inputs = self._get_inputs(docs, **kwargs)
-        prompt = self.llm_chain.prompt.format(**inputs)
-        return self.llm_chain.llm.get_num_tokens(prompt)
-
-    @property
-    def _chain_type(self) -> str:
-        return "stuff_documents_chain"
-
-
-class LoadingCallable(Protocol):
-    """Interface for loading the combine documents chain."""
-
-    def __call__(
-        self, llm: BaseLanguageModel, **kwargs: Any
-    ) -> BaseCombineDocumentsChain:
-        """Callable to load the combine documents chain."""
-
-
-def _load_stuff_chain(
-    llm: BaseLanguageModel,
-    prompt: Optional[BasePromptTemplate] = None,
-    document_variable_name: str = "context",
-    verbose: Optional[bool] = None,
-    callback_manager: Optional[BaseCallbackManager] = None,
-    callbacks: Callbacks = None,
-    **kwargs: Any,
-) -> StuffDocumentsChain:
-    _prompt = prompt or stuff_prompt.PROMPT_SELECTOR.get_prompt(llm)
-    llm_chain = LLMChain(
-        llm=llm,
-        prompt=_prompt,
-        verbose=verbose,
-        callback_manager=callback_manager,
-        callbacks=callbacks,
-    )
-    # TODO: document prompt
-    return StuffDocumentsChain(
-        llm_chain=llm_chain,
-        document_variable_name=document_variable_name,
-        verbose=verbose,
-        callback_manager=callback_manager,
-        **kwargs,
-    )
-
-
-def load_qa_chain(
-    llm: BaseLanguageModel,
-    chain_type: str = "stuff",
-    verbose: Optional[bool] = None,
-    callback_manager: Optional[BaseCallbackManager] = None,
-    **kwargs: Any,
-) -> BaseCombineDocumentsChain:
-    """Load question answering chain.
-
-    Args:
-        llm: Language Model to use in the chain.
-        chain_type: Type of document combining chain to use. Should be one of "stuff",
-            "map_reduce", "map_rerank", and "refine".
-        verbose: Whether chains should be run in verbose mode or not. Note that this
-            applies to all chains that make up the final chain.
-        callback_manager: Callback manager to use for the chain.
-
-    Returns:
-        A chain to use for question answering.
-    """
-    loader_mapping: Mapping[str, LoadingCallable] = {
-        "stuff": _load_stuff_chain,
-    }
-    if chain_type not in loader_mapping:
-        raise ValueError(
-            f"Got unsupported chain type: {chain_type}. "
-            f"Should be one of {loader_mapping.keys()}"
-        )
-    return loader_mapping[chain_type](
-        llm, verbose=verbose, callback_manager=callback_manager, **kwargs
-    )

apps/language_models/langchain/gpt4all_llm.py

@@ -1,380 +0,0 @@
-import inspect
-import os
-from functools import partial
-from typing import Dict, Any, Optional, List
-from langchain.callbacks.manager import CallbackManagerForLLMRun
-from pydantic import root_validator
-from langchain.llms import gpt4all
-from dotenv import dotenv_values
-
-from utils import FakeTokenizer
-
-
-def get_model_tokenizer_gpt4all(base_model, **kwargs):
-    # defaults (some of these are generation parameters, so need to be passed in at generation time)
-    model_kwargs = dict(
-        n_threads=os.cpu_count() // 2,
-        temp=kwargs.get("temperature", 0.2),
-        top_p=kwargs.get("top_p", 0.75),
-        top_k=kwargs.get("top_k", 40),
-        n_ctx=2048 - 256,
-    )
-    env_gpt4all_file = ".env_gpt4all"
-    model_kwargs.update(dotenv_values(env_gpt4all_file))
-    # make int or float if can to satisfy types for class
-    for k, v in model_kwargs.items():
-        try:
-            if float(v) == int(v):
-                model_kwargs[k] = int(v)
-            else:
-                model_kwargs[k] = float(v)
-        except:
-            pass
-
-    if base_model == "llama":
-        if "model_path_llama" not in model_kwargs:
-            raise ValueError("No model_path_llama in %s" % env_gpt4all_file)
-        model_path = model_kwargs.pop("model_path_llama")
-        # FIXME: GPT4All version of llama doesn't handle new quantization, so use llama_cpp_python
-        from llama_cpp import Llama
-
-        # llama sets some things at init model time, not generation time
-        func_names = list(inspect.signature(Llama.__init__).parameters)
-        model_kwargs = {
-            k: v for k, v in model_kwargs.items() if k in func_names
-        }
-        model_kwargs["n_ctx"] = int(model_kwargs["n_ctx"])
-        model = Llama(model_path=model_path, **model_kwargs)
-    elif base_model in "gpt4all_llama":
-        if (
-            "model_name_gpt4all_llama" not in model_kwargs
-            and "model_path_gpt4all_llama" not in model_kwargs
-        ):
-            raise ValueError(
-                "No model_name_gpt4all_llama or model_path_gpt4all_llama in %s"
-                % env_gpt4all_file
-            )
-        model_name = model_kwargs.pop("model_name_gpt4all_llama")
-        model_type = "llama"
-        from gpt4all import GPT4All as GPT4AllModel
-
-        model = GPT4AllModel(model_name=model_name, model_type=model_type)
-    elif base_model in "gptj":
-        if (
-            "model_name_gptj" not in model_kwargs
-            and "model_path_gptj" not in model_kwargs
-        ):
-            raise ValueError(
-                "No model_name_gpt4j or model_path_gpt4j in %s"
-                % env_gpt4all_file
-            )
-        model_name = model_kwargs.pop("model_name_gptj")
-        model_type = "gptj"
-        from gpt4all import GPT4All as GPT4AllModel
-
-        model = GPT4AllModel(model_name=model_name, model_type=model_type)
-    else:
-        raise ValueError("No such base_model %s" % base_model)
-    return model, FakeTokenizer(), "cpu"
-
-
-from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-
-
-class H2OStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Run on new LLM token. Only available when streaming is enabled."""
-        # streaming to std already occurs without this
-        # sys.stdout.write(token)
-        # sys.stdout.flush()
-        pass
-
-
-def get_model_kwargs(env_kwargs, default_kwargs, cls, exclude_list=[]):
-    # default from class
-    model_kwargs = {
-        k: v.default
-        for k, v in dict(inspect.signature(cls).parameters).items()
-        if k not in exclude_list
-    }
-    # from our defaults
-    model_kwargs.update(default_kwargs)
-    # from user defaults
-    model_kwargs.update(env_kwargs)
-    # ensure only valid keys
-    func_names = list(inspect.signature(cls).parameters)
-    model_kwargs = {k: v for k, v in model_kwargs.items() if k in func_names}
-    return model_kwargs
-
-
-def get_llm_gpt4all(
-    model_name,
-    model=None,
-    max_new_tokens=256,
-    temperature=0.1,
-    repetition_penalty=1.0,
-    top_k=40,
-    top_p=0.7,
-    streaming=False,
-    callbacks=None,
-    prompter=None,
-    verbose=False,
-):
-    assert prompter is not None
-    env_gpt4all_file = ".env_gpt4all"
-    env_kwargs = dotenv_values(env_gpt4all_file)
-    n_ctx = env_kwargs.pop("n_ctx", 2048 - max_new_tokens)
-    default_kwargs = dict(
-        context_erase=0.5,
-        n_batch=1,
-        n_ctx=n_ctx,
-        n_predict=max_new_tokens,
-        repeat_last_n=64 if repetition_penalty != 1.0 else 0,
-        repeat_penalty=repetition_penalty,
-        temp=temperature,
-        temperature=temperature,
-        top_k=top_k,
-        top_p=top_p,
-        use_mlock=True,
-        verbose=verbose,
-    )
-    if model_name == "llama":
-        cls = H2OLlamaCpp
-        model_path = (
-            env_kwargs.pop("model_path_llama") if model is None else model
-        )
-        model_kwargs = get_model_kwargs(
-            env_kwargs, default_kwargs, cls, exclude_list=["lc_kwargs"]
-        )
-        model_kwargs.update(
-            dict(
-                model_path=model_path,
-                callbacks=callbacks,
-                streaming=streaming,
-                prompter=prompter,
-            )
-        )
-        llm = cls(**model_kwargs)
-        llm.client.verbose = verbose
-    elif model_name == "gpt4all_llama":
-        cls = H2OGPT4All
-        model_path = (
-            env_kwargs.pop("model_path_gpt4all_llama")
-            if model is None
-            else model
-        )
-        model_kwargs = get_model_kwargs(
-            env_kwargs, default_kwargs, cls, exclude_list=["lc_kwargs"]
-        )
-        model_kwargs.update(
-            dict(
-                model=model_path,
-                backend="llama",
-                callbacks=callbacks,
-                streaming=streaming,
-                prompter=prompter,
-            )
-        )
-        llm = cls(**model_kwargs)
-    elif model_name == "gptj":
-        cls = H2OGPT4All
-        model_path = (
-            env_kwargs.pop("model_path_gptj") if model is None else model
-        )
-        model_kwargs = get_model_kwargs(
-            env_kwargs, default_kwargs, cls, exclude_list=["lc_kwargs"]
-        )
-        model_kwargs.update(
-            dict(
-                model=model_path,
-                backend="gptj",
-                callbacks=callbacks,
-                streaming=streaming,
-                prompter=prompter,
-            )
-        )
-        llm = cls(**model_kwargs)
-    else:
-        raise RuntimeError("No such model_name %s" % model_name)
-    return llm
-
-
-class H2OGPT4All(gpt4all.GPT4All):
-    model: Any
-    prompter: Any
-    """Path to the pre-trained GPT4All model file."""
-
-    @root_validator()
-    def validate_environment(cls, values: Dict) -> Dict:
-        """Validate that the python package exists in the environment."""
-        try:
-            if isinstance(values["model"], str):
-                from gpt4all import GPT4All as GPT4AllModel
-
-                full_path = values["model"]
-                model_path, delimiter, model_name = full_path.rpartition("/")
-                model_path += delimiter
-
-                values["client"] = GPT4AllModel(
-                    model_name=model_name,
-                    model_path=model_path or None,
-                    model_type=values["backend"],
-                    allow_download=False,
-                )
-                if values["n_threads"] is not None:
-                    # set n_threads
-                    values["client"].model.set_thread_count(
-                        values["n_threads"]
-                    )
-            else:
-                values["client"] = values["model"]
-            try:
-                values["backend"] = values["client"].model_type
-            except AttributeError:
-                # The below is for compatibility with GPT4All Python bindings <= 0.2.3.
-                values["backend"] = values["client"].model.model_type
-
-        except ImportError:
-            raise ValueError(
-                "Could not import gpt4all python package. "
-                "Please install it with `pip install gpt4all`."
-            )
-        return values
-
-    def _call(
-        self,
-        prompt: str,
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs,
-    ) -> str:
-        # Roughly 4 chars per token if natural language
-        prompt = prompt[-self.n_ctx * 4 :]
-
-        # use instruct prompting
-        data_point = dict(context="", instruction=prompt, input="")
-        prompt = self.prompter.generate_prompt(data_point)
-
-        verbose = False
-        if verbose:
-            print("_call prompt: %s" % prompt, flush=True)
-        # FIXME: GPT4ALl doesn't support yield during generate, so cannot support streaming except via itself to stdout
-        return super()._call(prompt, stop=stop, run_manager=run_manager)
-
-
-from langchain.llms import LlamaCpp
-
-
-class H2OLlamaCpp(LlamaCpp):
-    model_path: Any
-    prompter: Any
-    """Path to the pre-trained GPT4All model file."""
-
-    @root_validator()
-    def validate_environment(cls, values: Dict) -> Dict:
-        """Validate that llama-cpp-python library is installed."""
-        if isinstance(values["model_path"], str):
-            model_path = values["model_path"]
-            model_param_names = [
-                "lora_path",
-                "lora_base",
-                "n_ctx",
-                "n_parts",
-                "seed",
-                "f16_kv",
-                "logits_all",
-                "vocab_only",
-                "use_mlock",
-                "n_threads",
-                "n_batch",
-                "use_mmap",
-                "last_n_tokens_size",
-            ]
-            model_params = {k: values[k] for k in model_param_names}
-            # For backwards compatibility, only include if non-null.
-            if values["n_gpu_layers"] is not None:
-                model_params["n_gpu_layers"] = values["n_gpu_layers"]
-
-            try:
-                from llama_cpp import Llama
-
-                values["client"] = Llama(model_path, **model_params)
-            except ImportError:
-                raise ModuleNotFoundError(
-                    "Could not import llama-cpp-python library. "
-                    "Please install the llama-cpp-python library to "
-                    "use this embedding model: pip install llama-cpp-python"
-                )
-            except Exception as e:
-                raise ValueError(
-                    f"Could not load Llama model from path: {model_path}. "
-                    f"Received error {e}"
-                )
-        else:
-            values["client"] = values["model_path"]
-        return values
-
-    def _call(
-        self,
-        prompt: str,
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs,
-    ) -> str:
-        verbose = False
-        # tokenize twice, just to count tokens, since llama cpp python wrapper has no way to truncate
-        # still have to avoid crazy sizes, else hit llama_tokenize: too many tokens -- might still hit, not fatal
-        prompt = prompt[-self.n_ctx * 4 :]
-        prompt_tokens = self.client.tokenize(b" " + prompt.encode("utf-8"))
-        num_prompt_tokens = len(prompt_tokens)
-        if num_prompt_tokens > self.n_ctx:
-            # conservative by using int()
-            chars_per_token = int(len(prompt) / num_prompt_tokens)
-            prompt = prompt[-self.n_ctx * chars_per_token :]
-            if verbose:
-                print(
-                    "reducing tokens, assuming average of %s chars/token: %s"
-                    % chars_per_token,
-                    flush=True,
-                )
-                prompt_tokens2 = self.client.tokenize(
-                    b" " + prompt.encode("utf-8")
-                )
-                num_prompt_tokens2 = len(prompt_tokens2)
-                print(
-                    "reduced tokens from %d -> %d"
-                    % (num_prompt_tokens, num_prompt_tokens2),
-                    flush=True,
-                )
-
-        # use instruct prompting
-        data_point = dict(context="", instruction=prompt, input="")
-        prompt = self.prompter.generate_prompt(data_point)
-
-        if verbose:
-            print("_call prompt: %s" % prompt, flush=True)
-
-        if self.streaming:
-            text_callback = None
-            if run_manager:
-                text_callback = partial(
-                    run_manager.on_llm_new_token, verbose=self.verbose
-                )
-            # parent handler of streamer expects to see prompt first else output="" and lose if prompt=None in prompter
-            if text_callback:
-                text_callback(prompt)
-            text = ""
-            for token in self.stream(
-                prompt=prompt, stop=stop, run_manager=run_manager
-            ):
-                text_chunk = token["choices"][0]["text"]
-                # self.stream already calls text_callback
-                # if text_callback:
-                #    text_callback(text_chunk)
-                text += text_chunk
-            return text
-        else:
-            params = self._get_parameters(stop)
-            params = {**params, **kwargs}
-            result = self.client(prompt=prompt, **params)
-            return result["choices"][0]["text"]

apps/language_models/langchain/gradio_utils/grclient.py

@@ -1,93 +0,0 @@
-import traceback
-from typing import Callable
-import os
-
-from gradio_client.client import Job
-
-os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
-
-from gradio_client import Client
-
-
-class GradioClient(Client):
-    """
-    Parent class of gradio client
-    To handle automatically refreshing client if detect gradio server changed
-    """
-
-    def __init__(self, *args, **kwargs):
-        self.args = args
-        self.kwargs = kwargs
-        super().__init__(*args, **kwargs)
-        self.server_hash = self.get_server_hash()
-
-    def get_server_hash(self):
-        """
-        Get server hash using super without any refresh action triggered
-        Returns: git hash of gradio server
-        """
-        return super().submit(api_name="/system_hash").result()
-
-    def refresh_client_if_should(self):
-        # get current hash in order to update api_name -> fn_index map in case gradio server changed
-        # FIXME: Could add cli api as hash
-        server_hash = self.get_server_hash()
-        if self.server_hash != server_hash:
-            self.refresh_client()
-            self.server_hash = server_hash
-        else:
-            self.reset_session()
-
-    def refresh_client(self):
-        """
-        Ensure every client call is independent
-        Also ensure map between api_name and fn_index is updated in case server changed (e.g. restarted with new code)
-        Returns:
-        """
-        # need session hash to be new every time, to avoid "generator already executing"
-        self.reset_session()
-
-        client = Client(*self.args, **self.kwargs)
-        for k, v in client.__dict__.items():
-            setattr(self, k, v)
-
-    def submit(
-        self,
-        *args,
-        api_name: str | None = None,
-        fn_index: int | None = None,
-        result_callbacks: Callable | list[Callable] | None = None,
-    ) -> Job:
-        # Note predict calls submit
-        try:
-            self.refresh_client_if_should()
-            job = super().submit(*args, api_name=api_name, fn_index=fn_index)
-        except Exception as e:
-            print("Hit e=%s" % str(e), flush=True)
-            # force reconfig in case only that
-            self.refresh_client()
-            job = super().submit(*args, api_name=api_name, fn_index=fn_index)
-
-        # see if immediately failed
-        e = job.future._exception
-        if e is not None:
-            print(
-                "GR job failed: %s %s"
-                % (str(e), "".join(traceback.format_tb(e.__traceback__))),
-                flush=True,
-            )
-            # force reconfig in case only that
-            self.refresh_client()
-            job = super().submit(*args, api_name=api_name, fn_index=fn_index)
-            e2 = job.future._exception
-            if e2 is not None:
-                print(
-                    "GR job failed again: %s\n%s"
-                    % (
-                        str(e2),
-                        "".join(traceback.format_tb(e2.__traceback__)),
-                    ),
-                    flush=True,
-                )
-
-        return job

apps/language_models/langchain/h2oai_pipeline.py

@@ -1,765 +0,0 @@
-import os
-from apps.stable_diffusion.src.utils.utils import _compile_module
-from io import BytesIO
-import torch_mlir
-
-from stopping import get_stopping
-from prompter import Prompter, PromptType
-
-from transformers import TextGenerationPipeline
-from transformers.pipelines.text_generation import ReturnType
-from transformers.generation import (
-    GenerationConfig,
-    LogitsProcessorList,
-    StoppingCriteriaList,
-)
-import copy
-import torch
-from transformers import AutoConfig, AutoModelForCausalLM
-import gc
-from pathlib import Path
-from shark.shark_inference import SharkInference
-from shark.shark_downloader import download_public_file
-from shark.shark_importer import import_with_fx, save_mlir
-from apps.stable_diffusion.src import args
-
-# Brevitas
-from typing import List, Tuple
-from brevitas_examples.common.generative.quantize import quantize_model
-from brevitas_examples.llm.llm_quant.run_utils import get_model_impl
-
-
-# fmt: off
-def quant〇matmul_rhs_group_quant〡shape(lhs: List[int], rhs: List[int], rhs_scale: List[int], rhs_zero_point: List[int], rhs_bit_width: int, rhs_group_size: int) -> List[int]:
-    if len(lhs) == 3 and len(rhs) == 2:
-        return [lhs[0], lhs[1], rhs[0]]
-    elif len(lhs) == 2 and len(rhs) == 2:
-        return [lhs[0], rhs[0]]
-    else:
-        raise ValueError("Input shapes not supported.")
-
-
-def quant〇matmul_rhs_group_quant〡dtype(lhs_rank_dtype: Tuple[int, int], rhs_rank_dtype: Tuple[int, int], rhs_scale_rank_dtype: Tuple[int, int], rhs_zero_point_rank_dtype: Tuple[int, int], rhs_bit_width: int, rhs_group_size: int) -> int:
-    # output dtype is the dtype of the lhs float input
-    lhs_rank, lhs_dtype = lhs_rank_dtype
-    return lhs_dtype
-
-
-def quant〇matmul_rhs_group_quant〡has_value_semantics(lhs, rhs, rhs_scale, rhs_zero_point, rhs_bit_width, rhs_group_size) -> None:
-    return
-
-
-brevitas_matmul_rhs_group_quant_library = [
-    quant〇matmul_rhs_group_quant〡shape,
-    quant〇matmul_rhs_group_quant〡dtype,
-    quant〇matmul_rhs_group_quant〡has_value_semantics]
-# fmt: on
-
-global_device = "cuda"
-global_precision = "fp16"
-
-if not args.run_docuchat_web:
-    args.device = global_device
-    args.precision = global_precision
-tensor_device = "cpu" if args.device == "cpu" else "cuda"
-
-
-class H2OGPTModel(torch.nn.Module):
-    def __init__(self, device, precision):
-        super().__init__()
-        torch_dtype = (
-            torch.float32
-            if precision == "fp32" or device == "cpu"
-            else torch.float16
-        )
-        device_map = {"": "cpu"} if device == "cpu" else {"": 0}
-        model_kwargs = {
-            "local_files_only": False,
-            "torch_dtype": torch_dtype,
-            "resume_download": True,
-            "use_auth_token": False,
-            "trust_remote_code": True,
-            "offload_folder": "offline_folder",
-            "device_map": device_map,
-        }
-        config = AutoConfig.from_pretrained(
-            "h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v3",
-            use_auth_token=False,
-            trust_remote_code=True,
-            offload_folder="offline_folder",
-        )
-        self.model = AutoModelForCausalLM.from_pretrained(
-            "h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v3",
-            config=config,
-            **model_kwargs,
-        )
-        if precision in ["int4", "int8"]:
-            print("Applying weight quantization..")
-            weight_bit_width = 4 if precision == "int4" else 8
-            quantize_model(
-                self.model.transformer.h,
-                dtype=torch.float32,
-                weight_bit_width=weight_bit_width,
-                weight_param_method="stats",
-                weight_scale_precision="float_scale",
-                weight_quant_type="asym",
-                weight_quant_granularity="per_group",
-                weight_group_size=128,
-                quantize_weight_zero_point=False,
-            )
-            print("Weight quantization applied.")
-
-    def forward(self, input_ids, attention_mask):
-        input_dict = {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "past_key_values": None,
-            "use_cache": True,
-        }
-        output = self.model(
-            **input_dict,
-            return_dict=True,
-            output_attentions=False,
-            output_hidden_states=False,
-        )
-        return output.logits[:, -1, :]
-
-
-class H2OGPTSHARKModel(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        model_name = "h2ogpt_falcon_7b"
-        extended_model_name = (
-            model_name + "_" + args.precision + "_" + args.device
-        )
-        vmfb_path = Path(extended_model_name + ".vmfb")
-        mlir_path = Path(model_name + "_" + args.precision + ".mlir")
-        shark_module = None
-
-        need_to_compile = False
-        if not vmfb_path.exists():
-            need_to_compile = True
-            # Downloading VMFB from shark_tank
-            print("Trying to download pre-compiled vmfb from shark tank.")
-            download_public_file(
-                "gs://shark_tank/langchain/" + str(vmfb_path),
-                vmfb_path.absolute(),
-                single_file=True,
-            )
-            if vmfb_path.exists():
-                print(
-                    "Pre-compiled vmfb downloaded from shark tank successfully."
-                )
-                need_to_compile = False
-
-        if need_to_compile:
-            if not mlir_path.exists():
-                print("Trying to download pre-generated mlir from shark tank.")
-                # Downloading MLIR from shark_tank
-                download_public_file(
-                    "gs://shark_tank/langchain/" + str(mlir_path),
-                    mlir_path.absolute(),
-                    single_file=True,
-                )
-            if mlir_path.exists():
-                with open(mlir_path, "rb") as f:
-                    bytecode = f.read()
-            else:
-                # Generating the mlir
-                bytecode = self.get_bytecode(tensor_device, args.precision)
-
-            shark_module = SharkInference(
-                mlir_module=bytecode,
-                device=args.device,
-                mlir_dialect="linalg",
-            )
-            print(f"[DEBUG] generating vmfb.")
-            shark_module = _compile_module(
-                shark_module, extended_model_name, []
-            )
-            print("Saved newly generated vmfb.")
-
-        if shark_module is None:
-            if vmfb_path.exists():
-                print("Compiled vmfb found. Loading it from: ", vmfb_path)
-                shark_module = SharkInference(
-                    None, device=args.device, mlir_dialect="linalg"
-                )
-                shark_module.load_module(str(vmfb_path))
-                print("Compiled vmfb loaded successfully.")
-            else:
-                raise ValueError("Unable to download/generate a vmfb.")
-
-        self.model = shark_module
-
-    def get_bytecode(self, device, precision):
-        h2ogpt_model = H2OGPTModel(device, precision)
-
-        compilation_input_ids = torch.randint(
-            low=1, high=10000, size=(1, 400)
-        ).to(device=device)
-        compilation_attention_mask = torch.ones(1, 400, dtype=torch.int64).to(
-            device=device
-        )
-
-        h2ogptCompileInput = (
-            compilation_input_ids,
-            compilation_attention_mask,
-        )
-
-        print(f"[DEBUG] generating torchscript graph")
-        ts_graph = import_with_fx(
-            h2ogpt_model,
-            h2ogptCompileInput,
-            is_f16=False,
-            precision=precision,
-            f16_input_mask=[False, False],
-            mlir_type="torchscript",
-        )
-        del h2ogpt_model
-        del self.src_model
-
-        print(f"[DEBUG] generating torch mlir")
-        if precision in ["int4", "int8"]:
-            from torch_mlir.compiler_utils import (
-                run_pipeline_with_repro_report,
-            )
-
-            module = torch_mlir.compile(
-                ts_graph,
-                [*h2ogptCompileInput],
-                output_type=torch_mlir.OutputType.TORCH,
-                backend_legal_ops=["quant.matmul_rhs_group_quant"],
-                extra_library=brevitas_matmul_rhs_group_quant_library,
-                use_tracing=False,
-                verbose=False,
-            )
-            print(f"[DEBUG] converting torch to linalg")
-            run_pipeline_with_repro_report(
-                module,
-                "builtin.module(func.func(torch-unpack-quant-tensor),func.func(torch-convert-custom-quant-op),torch-backend-to-linalg-on-tensors-backend-pipeline)",
-                description="Lowering Torch Backend IR -> Linalg-on-Tensors Backend IR",
-            )
-        else:
-            module = torch_mlir.compile(
-                ts_graph,
-                [*h2ogptCompileInput],
-                torch_mlir.OutputType.LINALG_ON_TENSORS,
-                use_tracing=False,
-                verbose=False,
-            )
-        del ts_graph
-
-        print(f"[DEBUG] converting to bytecode")
-        bytecode_stream = BytesIO()
-        module.operation.write_bytecode(bytecode_stream)
-        bytecode = bytecode_stream.getvalue()
-        del module
-
-        bytecode = save_mlir(
-            bytecode,
-            model_name=f"h2ogpt_{precision}",
-            frontend="torch",
-        )
-        return bytecode
-
-    def forward(self, input_ids, attention_mask):
-        result = torch.from_numpy(
-            self.model(
-                "forward",
-                (input_ids.to(device="cpu"), attention_mask.to(device="cpu")),
-            )
-        ).to(device=tensor_device)
-        return result
-
-
-def decode_tokens(tokenizer, res_tokens):
-    for i in range(len(res_tokens)):
-        if type(res_tokens[i]) != int:
-            res_tokens[i] = int(res_tokens[i][0])
-
-    res_str = tokenizer.decode(res_tokens, skip_special_tokens=True)
-    return res_str
-
-
-def generate_token(h2ogpt_shark_model, model, tokenizer, **generate_kwargs):
-    del generate_kwargs["max_time"]
-    generate_kwargs["input_ids"] = generate_kwargs["input_ids"].to(
-        device=tensor_device
-    )
-    generate_kwargs["attention_mask"] = generate_kwargs["attention_mask"].to(
-        device=tensor_device
-    )
-    truncated_input_ids = []
-    stopping_criteria = generate_kwargs["stopping_criteria"]
-
-    generation_config_ = GenerationConfig.from_model_config(model.config)
-    generation_config = copy.deepcopy(generation_config_)
-    model_kwargs = generation_config.update(**generate_kwargs)
-
-    logits_processor = LogitsProcessorList()
-    stopping_criteria = (
-        stopping_criteria
-        if stopping_criteria is not None
-        else StoppingCriteriaList()
-    )
-
-    eos_token_id = generation_config.eos_token_id
-    generation_config.pad_token_id = eos_token_id
-
-    (
-        inputs_tensor,
-        model_input_name,
-        model_kwargs,
-    ) = model._prepare_model_inputs(
-        None, generation_config.bos_token_id, model_kwargs
-    )
-
-    model_kwargs["output_attentions"] = generation_config.output_attentions
-    model_kwargs[
-        "output_hidden_states"
-    ] = generation_config.output_hidden_states
-    model_kwargs["use_cache"] = generation_config.use_cache
-
-    input_ids = (
-        inputs_tensor
-        if model_input_name == "input_ids"
-        else model_kwargs.pop("input_ids")
-    )
-
-    input_ids_seq_length = input_ids.shape[-1]
-
-    generation_config.max_length = (
-        generation_config.max_new_tokens + input_ids_seq_length
-    )
-
-    logits_processor = model._get_logits_processor(
-        generation_config=generation_config,
-        input_ids_seq_length=input_ids_seq_length,
-        encoder_input_ids=inputs_tensor,
-        prefix_allowed_tokens_fn=None,
-        logits_processor=logits_processor,
-    )
-
-    stopping_criteria = model._get_stopping_criteria(
-        generation_config=generation_config,
-        stopping_criteria=stopping_criteria,
-    )
-
-    logits_warper = model._get_logits_warper(generation_config)
-
-    (
-        input_ids,
-        model_kwargs,
-    ) = model._expand_inputs_for_generation(
-        input_ids=input_ids,
-        expand_size=generation_config.num_return_sequences,  # 1
-        is_encoder_decoder=model.config.is_encoder_decoder,  # False
-        **model_kwargs,
-    )
-
-    if isinstance(eos_token_id, int):
-        eos_token_id = [eos_token_id]
-    eos_token_id_tensor = (
-        torch.tensor(eos_token_id).to(device=tensor_device)
-        if eos_token_id is not None
-        else None
-    )
-
-    pad_token_id = generation_config.pad_token_id
-    eos_token_id = eos_token_id
-
-    output_scores = generation_config.output_scores  # False
-    return_dict_in_generate = (
-        generation_config.return_dict_in_generate  # False
-    )
-
-    # init attention / hidden states / scores tuples
-    scores = () if (return_dict_in_generate and output_scores) else None
-
-    # keep track of which sequences are already finished
-    unfinished_sequences = torch.ones(
-        input_ids.shape[0],
-        dtype=torch.long,
-        device=input_ids.device,
-    )
-
-    timesRan = 0
-    import time
-
-    start = time.time()
-    print("\n")
-
-    res_tokens = []
-    while True:
-        model_inputs = model.prepare_inputs_for_generation(
-            input_ids, **model_kwargs
-        )
-
-        outputs = h2ogpt_shark_model.forward(
-            model_inputs["input_ids"], model_inputs["attention_mask"]
-        )
-
-        if args.precision == "fp16":
-            outputs = outputs.to(dtype=torch.float32)
-        next_token_logits = outputs
-
-        # pre-process distribution
-        next_token_scores = logits_processor(input_ids, next_token_logits)
-        next_token_scores = logits_warper(input_ids, next_token_scores)
-
-        # sample
-        probs = torch.nn.functional.softmax(next_token_scores, dim=-1)
-
-        next_token = torch.multinomial(probs, num_samples=1).squeeze(1)
-
-        # finished sentences should have their next token be a padding token
-        if eos_token_id is not None:
-            if pad_token_id is None:
-                raise ValueError(
-                    "If `eos_token_id` is defined, make sure that `pad_token_id` is defined."
-                )
-            next_token = next_token * unfinished_sequences + pad_token_id * (
-                1 - unfinished_sequences
-            )
-
-        input_ids = torch.cat([input_ids, next_token[:, None]], dim=-1)
-
-        model_kwargs["past_key_values"] = None
-        if "attention_mask" in model_kwargs:
-            attention_mask = model_kwargs["attention_mask"]
-            model_kwargs["attention_mask"] = torch.cat(
-                [
-                    attention_mask,
-                    attention_mask.new_ones((attention_mask.shape[0], 1)),
-                ],
-                dim=-1,
-            )
-
-        truncated_input_ids.append(input_ids[:, 0])
-        input_ids = input_ids[:, 1:]
-        model_kwargs["attention_mask"] = model_kwargs["attention_mask"][:, 1:]
-
-        new_word = tokenizer.decode(
-            next_token.cpu().numpy(),
-            add_special_tokens=False,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=True,
-        )
-
-        res_tokens.append(next_token)
-        if new_word == "<0x0A>":
-            print("\n", end="", flush=True)
-        else:
-            print(f"{new_word}", end=" ", flush=True)
-
-        part_str = decode_tokens(tokenizer, res_tokens)
-        yield part_str
-
-        # if eos_token was found in one sentence, set sentence to finished
-        if eos_token_id_tensor is not None:
-            unfinished_sequences = unfinished_sequences.mul(
-                next_token.tile(eos_token_id_tensor.shape[0], 1)
-                .ne(eos_token_id_tensor.unsqueeze(1))
-                .prod(dim=0)
-            )
-            # stop when each sentence is finished
-            if unfinished_sequences.max() == 0 or stopping_criteria(
-                input_ids, scores
-            ):
-                break
-        timesRan = timesRan + 1
-
-    end = time.time()
-    print(
-        "\n\nTime taken is {:.2f} seconds/token\n".format(
-            (end - start) / timesRan
-        )
-    )
-
-    torch.cuda.empty_cache()
-    gc.collect()
-
-    res_str = decode_tokens(tokenizer, res_tokens)
-    yield res_str
-
-
-def pad_or_truncate_inputs(
-    input_ids, attention_mask, max_padding_length=400, do_truncation=False
-):
-    inp_shape = input_ids.shape
-    if inp_shape[1] < max_padding_length:
-        # do padding
-        num_add_token = max_padding_length - inp_shape[1]
-        padded_input_ids = torch.cat(
-            [
-                torch.tensor([[11] * num_add_token]).to(device=tensor_device),
-                input_ids,
-            ],
-            dim=1,
-        )
-        padded_attention_mask = torch.cat(
-            [
-                torch.tensor([[0] * num_add_token]).to(device=tensor_device),
-                attention_mask,
-            ],
-            dim=1,
-        )
-        return padded_input_ids, padded_attention_mask
-    elif inp_shape[1] > max_padding_length or do_truncation:
-        # do truncation
-        num_remove_token = inp_shape[1] - max_padding_length
-        truncated_input_ids = input_ids[:, num_remove_token:]
-        truncated_attention_mask = attention_mask[:, num_remove_token:]
-        return truncated_input_ids, truncated_attention_mask
-    else:
-        return input_ids, attention_mask
-
-
-class H2OTextGenerationPipeline(TextGenerationPipeline):
-    def __init__(
-        self,
-        *args,
-        debug=False,
-        chat=False,
-        stream_output=False,
-        sanitize_bot_response=False,
-        use_prompter=True,
-        prompter=None,
-        prompt_type=None,
-        prompt_dict=None,
-        max_input_tokens=2048 - 256,
-        **kwargs,
-    ):
-        """
-        HF-like pipeline, but handle instruction prompting and stopping (for some models)
-        :param args:
-        :param debug:
-        :param chat:
-        :param stream_output:
-        :param sanitize_bot_response:
-        :param use_prompter: Whether to use prompter.  If pass prompt_type, will make prompter
-        :param prompter: prompter, can pass if have already
-        :param prompt_type: prompt_type, e.g. human_bot.  See prompt_type to model mapping in from prompter.py.
-                            If use_prompter, then will make prompter and use it.
-        :param prompt_dict: dict of get_prompt(, return_dict=True) for prompt_type=custom
-        :param max_input_tokens:
-        :param kwargs:
-        """
-        super().__init__(*args, **kwargs)
-        self.prompt_text = None
-        self.use_prompter = use_prompter
-        self.prompt_type = prompt_type
-        self.prompt_dict = prompt_dict
-        self.prompter = prompter
-        if self.use_prompter:
-            if self.prompter is not None:
-                assert self.prompter.prompt_type is not None
-            else:
-                self.prompter = Prompter(
-                    self.prompt_type,
-                    self.prompt_dict,
-                    debug=debug,
-                    chat=chat,
-                    stream_output=stream_output,
-                )
-            self.human = self.prompter.humanstr
-            self.bot = self.prompter.botstr
-            self.can_stop = True
-        else:
-            self.prompter = None
-            self.human = None
-            self.bot = None
-            self.can_stop = False
-        self.sanitize_bot_response = sanitize_bot_response
-        self.max_input_tokens = (
-            max_input_tokens  # not for generate, so ok that not kwargs
-        )
-
-    @staticmethod
-    def limit_prompt(prompt_text, tokenizer, max_prompt_length=None):
-        verbose = bool(int(os.getenv("VERBOSE_PIPELINE", "0")))
-
-        if hasattr(tokenizer, "model_max_length"):
-            # model_max_length only defined for generate.py, not raw use of h2oai_pipeline.py
-            model_max_length = tokenizer.model_max_length
-            if max_prompt_length is not None:
-                model_max_length = min(model_max_length, max_prompt_length)
-            # cut at some upper likely limit to avoid excessive tokenization etc
-            # upper bound of 10 chars/token, e.g. special chars sometimes are long
-            if len(prompt_text) > model_max_length * 10:
-                len0 = len(prompt_text)
-                prompt_text = prompt_text[-model_max_length * 10 :]
-                if verbose:
-                    print(
-                        "Cut of input: %s -> %s" % (len0, len(prompt_text)),
-                        flush=True,
-                    )
-        else:
-            # unknown
-            model_max_length = None
-
-        num_prompt_tokens = None
-        if model_max_length is not None:
-            # can't wait for "hole" if not plain prompt_type, since would lose prefix like <human>:
-            # For https://github.com/h2oai/h2ogpt/issues/192
-            for trial in range(0, 3):
-                prompt_tokens = tokenizer(prompt_text)["input_ids"]
-                num_prompt_tokens = len(prompt_tokens)
-                if num_prompt_tokens > model_max_length:
-                    # conservative by using int()
-                    chars_per_token = int(len(prompt_text) / num_prompt_tokens)
-                    # keep tail, where question is if using langchain
-                    prompt_text = prompt_text[
-                        -model_max_length * chars_per_token :
-                    ]
-                    if verbose:
-                        print(
-                            "reducing %s tokens, assuming average of %s chars/token for %s characters"
-                            % (
-                                num_prompt_tokens,
-                                chars_per_token,
-                                len(prompt_text),
-                            ),
-                            flush=True,
-                        )
-                else:
-                    if verbose:
-                        print(
-                            "using %s tokens with %s chars"
-                            % (num_prompt_tokens, len(prompt_text)),
-                            flush=True,
-                        )
-                    break
-
-        return prompt_text, num_prompt_tokens
-
-    def preprocess(
-        self,
-        prompt_text,
-        prefix="",
-        handle_long_generation=None,
-        **generate_kwargs,
-    ):
-        (
-            prompt_text,
-            num_prompt_tokens,
-        ) = H2OTextGenerationPipeline.limit_prompt(prompt_text, self.tokenizer)
-
-        data_point = dict(context="", instruction=prompt_text, input="")
-        if self.prompter is not None:
-            prompt_text = self.prompter.generate_prompt(data_point)
-        self.prompt_text = prompt_text
-        if handle_long_generation is None:
-            # forces truncation of inputs to avoid critical failure
-            handle_long_generation = None  # disable with new approaches
-        return super().preprocess(
-            prompt_text,
-            prefix=prefix,
-            handle_long_generation=handle_long_generation,
-            **generate_kwargs,
-        )
-
-    def postprocess(
-        self,
-        model_outputs,
-        return_type=ReturnType.FULL_TEXT,
-        clean_up_tokenization_spaces=True,
-    ):
-        records = super().postprocess(
-            model_outputs,
-            return_type=return_type,
-            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
-        )
-        for rec in records:
-            if self.use_prompter:
-                outputs = rec["generated_text"]
-                outputs = self.prompter.get_response(
-                    outputs,
-                    prompt=self.prompt_text,
-                    sanitize_bot_response=self.sanitize_bot_response,
-                )
-            elif self.bot and self.human:
-                outputs = (
-                    rec["generated_text"]
-                    .split(self.bot)[1]
-                    .split(self.human)[0]
-                )
-            else:
-                outputs = rec["generated_text"]
-            rec["generated_text"] = outputs
-            print(
-                "prompt: %s\noutputs: %s\n\n" % (self.prompt_text, outputs),
-                flush=True,
-            )
-        return records
-
-    def _forward(self, model_inputs, **generate_kwargs):
-        if self.can_stop:
-            stopping_criteria = get_stopping(
-                self.prompt_type,
-                self.prompt_dict,
-                self.tokenizer,
-                self.device,
-                human=self.human,
-                bot=self.bot,
-                model_max_length=self.tokenizer.model_max_length,
-            )
-            generate_kwargs["stopping_criteria"] = stopping_criteria
-        # return super()._forward(model_inputs, **generate_kwargs)
-        return self.__forward(model_inputs, **generate_kwargs)
-
-    # FIXME: Copy-paste of original _forward, but removed copy.deepcopy()
-    # FIXME: https://github.com/h2oai/h2ogpt/issues/172
-    def __forward(self, model_inputs, **generate_kwargs):
-        input_ids = model_inputs["input_ids"]
-        attention_mask = model_inputs.get("attention_mask", None)
-        # Allow empty prompts
-        if input_ids.shape[1] == 0:
-            input_ids = None
-            attention_mask = None
-            in_b = 1
-        else:
-            in_b = input_ids.shape[0]
-        prompt_text = model_inputs.pop("prompt_text")
-
-        ## If there is a prefix, we may need to adjust the generation length. Do so without permanently modifying
-        ## generate_kwargs, as some of the parameterization may come from the initialization of the pipeline.
-        # generate_kwargs = copy.deepcopy(generate_kwargs)
-        prefix_length = generate_kwargs.pop("prefix_length", 0)
-        if prefix_length > 0:
-            has_max_new_tokens = "max_new_tokens" in generate_kwargs or (
-                "generation_config" in generate_kwargs
-                and generate_kwargs["generation_config"].max_new_tokens
-                is not None
-            )
-            if not has_max_new_tokens:
-                generate_kwargs["max_length"] = (
-                    generate_kwargs.get("max_length")
-                    or self.model.config.max_length
-                )
-                generate_kwargs["max_length"] += prefix_length
-            has_min_new_tokens = "min_new_tokens" in generate_kwargs or (
-                "generation_config" in generate_kwargs
-                and generate_kwargs["generation_config"].min_new_tokens
-                is not None
-            )
-            if not has_min_new_tokens and "min_length" in generate_kwargs:
-                generate_kwargs["min_length"] += prefix_length
-
-        # BS x SL
-        # pad or truncate the input_ids and attention_mask
-        max_padding_length = 400
-        input_ids, attention_mask = pad_or_truncate_inputs(
-            input_ids, attention_mask, max_padding_length=max_padding_length
-        )
-
-        return_dict = {
-            "model": self.model,
-            "tokenizer": self.tokenizer,
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "attention_mask": attention_mask,
-        }
-        return_dict = {**return_dict, **generate_kwargs}
-        return return_dict

apps/language_models/langchain/image_captions.py

@@ -1,247 +0,0 @@
-"""
-Based upon ImageCaptionLoader in LangChain version: langchain/document_loaders/image_captions.py
-But accepts preloaded model to avoid slowness in use and CUDA forking issues
-
-Loader that loads image captions
-By default, the loader utilizes the pre-trained BLIP image captioning model.
-https://huggingface.co/Salesforce/blip-image-captioning-base
-
-"""
-from typing import List, Union, Any, Tuple
-
-import requests
-from langchain.docstore.document import Document
-from langchain.document_loaders import ImageCaptionLoader
-
-from utils import get_device, NullContext
-
-import pkg_resources
-
-try:
-    assert pkg_resources.get_distribution("bitsandbytes") is not None
-    have_bitsandbytes = True
-except (pkg_resources.DistributionNotFound, AssertionError):
-    have_bitsandbytes = False
-
-
-class H2OImageCaptionLoader(ImageCaptionLoader):
-    """Loader that loads the captions of an image"""
-
-    def __init__(
-        self,
-        path_images: Union[str, List[str]] = None,
-        blip_processor: str = None,
-        blip_model: str = None,
-        caption_gpu=True,
-        load_in_8bit=True,
-        # True doesn't seem to work, even though https://huggingface.co/Salesforce/blip2-flan-t5-xxl#in-8-bit-precision-int8
-        load_half=False,
-        load_gptq="",
-        use_safetensors=False,
-        min_new_tokens=20,
-        max_tokens=50,
-    ):
-        if blip_model is None or blip_model is None:
-            blip_processor = "Salesforce/blip-image-captioning-base"
-            blip_model = "Salesforce/blip-image-captioning-base"
-
-        super().__init__(path_images, blip_processor, blip_model)
-        self.blip_processor = blip_processor
-        self.blip_model = blip_model
-        self.processor = None
-        self.model = None
-        self.caption_gpu = caption_gpu
-        self.context_class = NullContext
-        self.device = "cpu"
-        self.load_in_8bit = (
-            load_in_8bit and have_bitsandbytes
-        )  # only for blip2
-        self.load_half = load_half
-        self.load_gptq = load_gptq
-        self.use_safetensors = use_safetensors
-        self.gpu_id = "auto"
-        # default prompt
-        self.prompt = "image of"
-        self.min_new_tokens = min_new_tokens
-        self.max_tokens = max_tokens
-
-    def set_context(self):
-        if get_device() == "cuda" and self.caption_gpu:
-            import torch
-
-            n_gpus = (
-                torch.cuda.device_count() if torch.cuda.is_available else 0
-            )
-            if n_gpus > 0:
-                self.context_class = torch.device
-                self.device = "cuda"
-
-    def load_model(self):
-        try:
-            import transformers
-        except ImportError:
-            raise ValueError(
-                "`transformers` package not found, please install with "
-                "`pip install transformers`."
-            )
-        self.set_context()
-        if self.caption_gpu:
-            if self.gpu_id == "auto":
-                # blip2 has issues with multi-GPU.  Error says need to somehow set language model in device map
-                # device_map = 'auto'
-                device_map = {"": 0}
-            else:
-                if self.device == "cuda":
-                    device_map = {"": self.gpu_id}
-                else:
-                    device_map = {"": "cpu"}
-        else:
-            device_map = {"": "cpu"}
-        import torch
-
-        with torch.no_grad():
-            with self.context_class(self.device):
-                context_class_cast = (
-                    NullContext if self.device == "cpu" else torch.autocast
-                )
-                with context_class_cast(self.device):
-                    if "blip2" in self.blip_processor.lower():
-                        from transformers import (
-                            Blip2Processor,
-                            Blip2ForConditionalGeneration,
-                        )
-
-                        if self.load_half and not self.load_in_8bit:
-                            self.processor = Blip2Processor.from_pretrained(
-                                self.blip_processor, device_map=device_map
-                            ).half()
-                            self.model = (
-                                Blip2ForConditionalGeneration.from_pretrained(
-                                    self.blip_model, device_map=device_map
-                                ).half()
-                            )
-                        else:
-                            self.processor = Blip2Processor.from_pretrained(
-                                self.blip_processor,
-                                load_in_8bit=self.load_in_8bit,
-                                device_map=device_map,
-                            )
-                            self.model = (
-                                Blip2ForConditionalGeneration.from_pretrained(
-                                    self.blip_model,
-                                    load_in_8bit=self.load_in_8bit,
-                                    device_map=device_map,
-                                )
-                            )
-                    else:
-                        from transformers import (
-                            BlipForConditionalGeneration,
-                            BlipProcessor,
-                        )
-
-                        self.load_half = False  # not supported
-                        if self.caption_gpu:
-                            if device_map == "auto":
-                                # Blip doesn't support device_map='auto'
-                                if self.device == "cuda":
-                                    if self.gpu_id == "auto":
-                                        device_map = {"": 0}
-                                    else:
-                                        device_map = {"": self.gpu_id}
-                                else:
-                                    device_map = {"": "cpu"}
-                        else:
-                            device_map = {"": "cpu"}
-                        self.processor = BlipProcessor.from_pretrained(
-                            self.blip_processor, device_map=device_map
-                        )
-                        self.model = (
-                            BlipForConditionalGeneration.from_pretrained(
-                                self.blip_model, device_map=device_map
-                            )
-                        )
-        return self
-
-    def set_image_paths(self, path_images: Union[str, List[str]]):
-        """
-        Load from a list of image files
-        """
-        if isinstance(path_images, str):
-            self.image_paths = [path_images]
-        else:
-            self.image_paths = path_images
-
-    def load(self, prompt=None) -> List[Document]:
-        if self.processor is None or self.model is None:
-            self.load_model()
-        results = []
-        for path_image in self.image_paths:
-            caption, metadata = self._get_captions_and_metadata(
-                model=self.model,
-                processor=self.processor,
-                path_image=path_image,
-                prompt=prompt,
-            )
-            doc = Document(page_content=caption, metadata=metadata)
-            results.append(doc)
-
-        return results
-
-    def _get_captions_and_metadata(
-        self, model: Any, processor: Any, path_image: str, prompt=None
-    ) -> Tuple[str, dict]:
-        """
-        Helper function for getting the captions and metadata of an image
-        """
-        if prompt is None:
-            prompt = self.prompt
-        try:
-            from PIL import Image
-        except ImportError:
-            raise ValueError(
-                "`PIL` package not found, please install with `pip install pillow`"
-            )
-
-        try:
-            if path_image.startswith("http://") or path_image.startswith(
-                "https://"
-            ):
-                image = Image.open(
-                    requests.get(path_image, stream=True).raw
-                ).convert("RGB")
-            else:
-                image = Image.open(path_image).convert("RGB")
-        except Exception:
-            raise ValueError(f"Could not get image data for {path_image}")
-
-        import torch
-
-        with torch.no_grad():
-            with self.context_class(self.device):
-                context_class_cast = (
-                    NullContext if self.device == "cpu" else torch.autocast
-                )
-                with context_class_cast(self.device):
-                    if self.load_half:
-                        inputs = processor(
-                            image, prompt, return_tensors="pt"
-                        ).half()
-                    else:
-                        inputs = processor(image, prompt, return_tensors="pt")
-                    min_length = len(prompt) // 4 + self.min_new_tokens
-                    self.max_tokens = max(self.max_tokens, min_length)
-                    output = model.generate(
-                        **inputs,
-                        min_length=min_length,
-                        max_length=self.max_tokens,
-                    )
-
-                    caption: str = processor.decode(
-                        output[0], skip_special_tokens=True
-                    )
-                    prompti = caption.find(prompt)
-                    if prompti >= 0:
-                        caption = caption[prompti + len(prompt) :]
-                    metadata: dict = {"image_path": path_image}
-
-        return caption, metadata

apps/language_models/langchain/langchain_requirements.txt

@@ -1,120 +0,0 @@
-# for generate (gradio server) and finetune
-datasets==2.13.0
-sentencepiece==0.1.99
-huggingface_hub==0.16.4
-appdirs==1.4.4
-fire==0.5.0
-docutils==0.20.1
-evaluate==0.4.0
-rouge_score==0.1.2
-sacrebleu==2.3.1
-scikit-learn==1.2.2
-alt-profanity-check==1.2.2
-better-profanity==0.7.0
-numpy==1.24.3
-pandas==2.0.2
-matplotlib==3.7.1
-loralib==0.1.1
-bitsandbytes==0.39.0
-accelerate==0.20.3
-peft==0.4.0
-# 4.31.0+ breaks load_in_8bit=True (https://github.com/huggingface/transformers/issues/25026)
-transformers==4.30.2
-tokenizers==0.13.3
-APScheduler==3.10.1
-
-# optional for generate
-pynvml==11.5.0
-psutil==5.9.5
-boto3==1.26.101
-botocore==1.29.101
-
-# optional for finetune
-tensorboard==2.13.0
-neptune==1.2.0
-
-# for gradio client
-gradio_client==0.2.10
-beautifulsoup4==4.12.2
-markdown==3.4.3
-
-# data and testing
-pytest==7.2.2
-pytest-xdist==3.2.1
-nltk==3.8.1
-textstat==0.7.3
-# pandoc==2.3
-pypandoc==1.11; sys_platform == "darwin" and platform_machine == "arm64"
-pypandoc_binary==1.11; platform_machine == "x86_64"
-pypandoc_binary==1.11; sys_platform == "win32"
-openpyxl==3.1.2
-lm_dataformat==0.0.20
-bioc==2.0
-
-# falcon
-einops==0.6.1
-instructorembedding==1.0.1
-
-# for gpt4all .env file, but avoid worrying about imports
-python-dotenv==1.0.0
-
-text-generation==0.6.0
-# for tokenization when don't have HF tokenizer
-tiktoken==0.4.0
-# optional: for OpenAI endpoint or embeddings (requires key)
-openai==0.27.8
-
-# optional for chat with PDF
-langchain==0.0.329
-pypdf==3.17.0
-# avoid textract, requires old six
-#textract==1.6.5
-
-# for HF embeddings
-sentence_transformers==2.2.2
-
-# local vector db
-chromadb==0.3.25
-# server vector db
-#pymilvus==2.2.8
-
-# weak url support, if can't install opencv etc. If comment-in this one, then comment-out unstructured[local-inference]==0.6.6
-# unstructured==0.8.1
-
-# strong support for images
-# Requires on Ubuntu: sudo apt-get install libmagic-dev poppler-utils tesseract-ocr libtesseract-dev libreoffice
-unstructured[local-inference]==0.7.4
-#pdf2image==1.16.3
-#pytesseract==0.3.10
-pillow
-
-pdfminer.six==20221105
-urllib3
-requests_file
-
-#pdf2image==1.16.3
-#pytesseract==0.3.10
-tabulate==0.9.0
-# FYI pandoc already part of requirements.txt
-
-# JSONLoader, but makes some trouble for some users
-# jq==1.4.1
-
-# to check licenses
-# Run: pip-licenses|grep -v 'BSD\|Apache\|MIT'
-pip-licenses==4.3.0
-
-# weaviate vector db
-weaviate-client==3.22.1
-
-gpt4all==1.0.5
-llama-cpp-python==0.1.73
-
-arxiv==1.4.8
-pymupdf==1.22.5 # AGPL license
-# extract-msg==0.41.1  # GPL3
-
-# sometimes unstructured fails, these work in those cases.  See https://github.com/h2oai/h2ogpt/issues/320
-playwright==1.36.0
-# requires Chrome binary to be in path
-selenium==4.10.0

apps/language_models/langchain/llama_flash_attn_monkey_patch.py

@@ -1,124 +0,0 @@
-from typing import List, Optional, Tuple
-
-import torch
-
-import transformers
-from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
-
-from einops import rearrange
-
-from flash_attn.flash_attn_interface import flash_attn_unpadded_qkvpacked_func
-from flash_attn.bert_padding import unpad_input, pad_input
-
-
-def forward(
-    self,
-    hidden_states: torch.Tensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.Tensor] = None,
-    past_key_value: Optional[Tuple[torch.Tensor]] = None,
-    output_attentions: bool = False,
-    use_cache: bool = False,
-) -> Tuple[
-    torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]
-]:
-    """Input shape: Batch x Time x Channel
-    attention_mask: [bsz, q_len]
-    """
-    bsz, q_len, _ = hidden_states.size()
-
-    query_states = (
-        self.q_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    key_states = (
-        self.k_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    value_states = (
-        self.v_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    # [bsz, q_len, nh, hd]
-    # [bsz, nh, q_len, hd]
-
-    kv_seq_len = key_states.shape[-2]
-    assert past_key_value is None, "past_key_value is not supported"
-
-    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-    query_states, key_states = apply_rotary_pos_emb(
-        query_states, key_states, cos, sin, position_ids
-    )
-    # [bsz, nh, t, hd]
-    assert not output_attentions, "output_attentions is not supported"
-    assert not use_cache, "use_cache is not supported"
-
-    # Flash attention codes from
-    # https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attention.py
-
-    # transform the data into the format required by flash attention
-    qkv = torch.stack(
-        [query_states, key_states, value_states], dim=2
-    )  # [bsz, nh, 3, q_len, hd]
-    qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
-    # We have disabled _prepare_decoder_attention_mask in LlamaModel
-    # the attention_mask should be the same as the key_padding_mask
-    key_padding_mask = attention_mask
-
-    if key_padding_mask is None:
-        qkv = rearrange(qkv, "b s ... -> (b s) ...")
-        max_s = q_len
-        cu_q_lens = torch.arange(
-            0,
-            (bsz + 1) * q_len,
-            step=q_len,
-            dtype=torch.int32,
-            device=qkv.device,
-        )
-        output = flash_attn_unpadded_qkvpacked_func(
-            qkv, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
-        )
-        output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
-    else:
-        nheads = qkv.shape[-2]
-        x = rearrange(qkv, "b s three h d -> b s (three h d)")
-        x_unpad, indices, cu_q_lens, max_s = unpad_input(x, key_padding_mask)
-        x_unpad = rearrange(
-            x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads
-        )
-        output_unpad = flash_attn_unpadded_qkvpacked_func(
-            x_unpad, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
-        )
-        output = rearrange(
-            pad_input(
-                rearrange(output_unpad, "nnz h d -> nnz (h d)"),
-                indices,
-                bsz,
-                q_len,
-            ),
-            "b s (h d) -> b s h d",
-            h=nheads,
-        )
-    return self.o_proj(rearrange(output, "b s h d -> b s (h d)")), None, None
-
-
-# Disable the transformation of the attention mask in LlamaModel as the flash attention
-# requires the attention mask to be the same as the key_padding_mask
-def _prepare_decoder_attention_mask(
-    self, attention_mask, input_shape, inputs_embeds, past_key_values_length
-):
-    # [bsz, seq_len]
-    return attention_mask
-
-
-def replace_llama_attn_with_flash_attn():
-    print(
-        "Replacing original LLaMa attention with flash attention", flush=True
-    )
-    transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (
-        _prepare_decoder_attention_mask
-    )
-    transformers.models.llama.modeling_llama.LlamaAttention.forward = forward

apps/language_models/langchain/loaders.py

@@ -1,109 +0,0 @@
-import functools
-
-
-def get_loaders(model_name, reward_type, llama_type=None, load_gptq=""):
-    # NOTE: Some models need specific new prompt_type
-    # E.g. t5_xxl_true_nli_mixture has input format: "premise: PREMISE_TEXT hypothesis: HYPOTHESIS_TEXT".)
-    if load_gptq:
-        from transformers import AutoTokenizer
-        from auto_gptq import AutoGPTQForCausalLM
-
-        use_triton = False
-        functools.partial(
-            AutoGPTQForCausalLM.from_quantized,
-            quantize_config=None,
-            use_triton=use_triton,
-        )
-        return AutoGPTQForCausalLM.from_quantized, AutoTokenizer
-    if llama_type is None:
-        llama_type = "llama" in model_name.lower()
-    if llama_type:
-        from transformers import LlamaForCausalLM, LlamaTokenizer
-
-        return LlamaForCausalLM.from_pretrained, LlamaTokenizer
-    elif "distilgpt2" in model_name.lower():
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-
-        return AutoModelForCausalLM.from_pretrained, AutoTokenizer
-    elif "gpt2" in model_name.lower():
-        from transformers import GPT2LMHeadModel, GPT2Tokenizer
-
-        return GPT2LMHeadModel.from_pretrained, GPT2Tokenizer
-    elif "mbart-" in model_name.lower():
-        from transformers import (
-            MBartForConditionalGeneration,
-            MBart50TokenizerFast,
-        )
-
-        return (
-            MBartForConditionalGeneration.from_pretrained,
-            MBart50TokenizerFast,
-        )
-    elif (
-        "t5" == model_name.lower()
-        or "t5-" in model_name.lower()
-        or "flan-" in model_name.lower()
-    ):
-        from transformers import AutoTokenizer, T5ForConditionalGeneration
-
-        return T5ForConditionalGeneration.from_pretrained, AutoTokenizer
-    elif "bigbird" in model_name:
-        from transformers import (
-            BigBirdPegasusForConditionalGeneration,
-            AutoTokenizer,
-        )
-
-        return (
-            BigBirdPegasusForConditionalGeneration.from_pretrained,
-            AutoTokenizer,
-        )
-    elif (
-        "bart-large-cnn-samsum" in model_name
-        or "flan-t5-base-samsum" in model_name
-    ):
-        from transformers import pipeline
-
-        return pipeline, "summarization"
-    elif (
-        reward_type
-        or "OpenAssistant/reward-model".lower() in model_name.lower()
-    ):
-        from transformers import (
-            AutoModelForSequenceClassification,
-            AutoTokenizer,
-        )
-
-        return (
-            AutoModelForSequenceClassification.from_pretrained,
-            AutoTokenizer,
-        )
-    else:
-        from transformers import AutoTokenizer, AutoModelForCausalLM
-
-        model_loader = AutoModelForCausalLM
-        tokenizer_loader = AutoTokenizer
-        return model_loader.from_pretrained, tokenizer_loader
-
-
-def get_tokenizer(
-    tokenizer_loader,
-    tokenizer_base_model,
-    local_files_only,
-    resume_download,
-    use_auth_token,
-):
-    tokenizer = tokenizer_loader.from_pretrained(
-        tokenizer_base_model,
-        local_files_only=local_files_only,
-        resume_download=resume_download,
-        use_auth_token=use_auth_token,
-        padding_side="left",
-    )
-
-    tokenizer.pad_token_id = 0  # different from the eos token
-    # when generating, we will use the logits of right-most token to predict the next token
-    # so the padding should be on the left,
-    # e.g. see: https://huggingface.co/transformers/v4.11.3/model_doc/t5.html#inference
-    tokenizer.padding_side = "left"  # Allow batched inference
-
-    return tokenizer

apps/language_models/langchain/make_db.py

@@ -1,203 +0,0 @@
-import os
-
-from gpt_langchain import (
-    path_to_docs,
-    get_some_dbs_from_hf,
-    all_db_zips,
-    some_db_zips,
-    create_or_update_db,
-)
-from utils import get_ngpus_vis
-
-
-def glob_to_db(
-    user_path,
-    chunk=True,
-    chunk_size=512,
-    verbose=False,
-    fail_any_exception=False,
-    n_jobs=-1,
-    url=None,
-    enable_captions=True,
-    captions_model=None,
-    caption_loader=None,
-    enable_ocr=False,
-):
-    sources1 = path_to_docs(
-        user_path,
-        verbose=verbose,
-        fail_any_exception=fail_any_exception,
-        n_jobs=n_jobs,
-        chunk=chunk,
-        chunk_size=chunk_size,
-        url=url,
-        enable_captions=enable_captions,
-        captions_model=captions_model,
-        caption_loader=caption_loader,
-        enable_ocr=enable_ocr,
-    )
-    return sources1
-
-
-def make_db_main(
-    use_openai_embedding: bool = False,
-    hf_embedding_model: str = None,
-    persist_directory: str = "db_dir_UserData",
-    user_path: str = "user_path",
-    url: str = None,
-    add_if_exists: bool = True,
-    collection_name: str = "UserData",
-    verbose: bool = False,
-    chunk: bool = True,
-    chunk_size: int = 512,
-    fail_any_exception: bool = False,
-    download_all: bool = False,
-    download_some: bool = False,
-    download_one: str = None,
-    download_dest: str = "./",
-    n_jobs: int = -1,
-    enable_captions: bool = True,
-    captions_model: str = "Salesforce/blip-image-captioning-base",
-    pre_load_caption_model: bool = False,
-    caption_gpu: bool = True,
-    enable_ocr: bool = False,
-    db_type: str = "chroma",
-):
-    """
-    # To make UserData db for generate.py, put pdfs, etc. into path user_path and run:
-    python make_db.py
-
-    # once db is made, can use in generate.py like:
-
-    python generate.py --base_model=h2oai/h2ogpt-oig-oasst1-512-6_9b --langchain_mode=UserData
-
-    or zip-up the db_dir_UserData and share:
-
-    zip -r db_dir_UserData.zip db_dir_UserData
-
-    # To get all db files (except large wiki_full) do:
-    python make_db.py --download_some=True
-
-    # To get a single db file from HF:
-    python make_db.py --download_one=db_dir_DriverlessAI_docs.zip
-
-    :param use_openai_embedding: Whether to use OpenAI embedding
-    :param hf_embedding_model: HF embedding model to use. Like generate.py, uses 'hkunlp/instructor-large' if have GPUs, else "sentence-transformers/all-MiniLM-L6-v2"
-    :param persist_directory: where to persist db
-    :param user_path: where to pull documents from (None means url is not None.  If url is not None, this is ignored.)
-    :param url: url to generate documents from (None means user_path is not None)
-    :param add_if_exists: Add to db if already exists, but will not add duplicate sources
-    :param collection_name: Collection name for new db if not adding
-    :param verbose: whether to show verbose messages
-    :param chunk: whether to chunk data
-    :param chunk_size: chunk size for chunking
-    :param fail_any_exception: whether to fail if any exception hit during ingestion of files
-    :param download_all: whether to download all (including 23GB Wikipedia) example databases from h2o.ai HF
-    :param download_some: whether to download some small example databases from h2o.ai HF
-    :param download_one: whether to download one chosen example databases from h2o.ai HF
-    :param download_dest: Destination for downloads
-    :param n_jobs: Number of cores to use for ingesting multiple files
-    :param enable_captions: Whether to enable captions on images
-    :param captions_model: See generate.py
-    :param pre_load_caption_model: See generate.py
-    :param caption_gpu: Caption images on GPU if present
-    :param enable_ocr: Whether to enable OCR on images
-    :param db_type: Type of db to create. Currently only 'chroma' and 'weaviate' is supported.
-    :return: None
-    """
-    db = None
-
-    # match behavior of main() in generate.py for non-HF case
-    n_gpus = get_ngpus_vis()
-    if n_gpus == 0:
-        if hf_embedding_model is None:
-            # if no GPUs, use simpler embedding model to avoid cost in time
-            hf_embedding_model = "sentence-transformers/all-MiniLM-L6-v2"
-    else:
-        if hf_embedding_model is None:
-            # if still None, then set default
-            hf_embedding_model = "hkunlp/instructor-large"
-
-    if download_all:
-        print("Downloading all (and unzipping): %s" % all_db_zips, flush=True)
-        get_some_dbs_from_hf(download_dest, db_zips=all_db_zips)
-        if verbose:
-            print("DONE", flush=True)
-        return db, collection_name
-    elif download_some:
-        print(
-            "Downloading some (and unzipping): %s" % some_db_zips, flush=True
-        )
-        get_some_dbs_from_hf(download_dest, db_zips=some_db_zips)
-        if verbose:
-            print("DONE", flush=True)
-        return db, collection_name
-    elif download_one:
-        print("Downloading %s (and unzipping)" % download_one, flush=True)
-        get_some_dbs_from_hf(
-            download_dest, db_zips=[[download_one, "", "Unknown License"]]
-        )
-        if verbose:
-            print("DONE", flush=True)
-        return db, collection_name
-
-    if enable_captions and pre_load_caption_model:
-        # preload, else can be too slow or if on GPU have cuda context issues
-        # Inside ingestion, this will disable parallel loading of multiple other kinds of docs
-        # However, if have many images, all those images will be handled more quickly by preloaded model on GPU
-        from image_captions import H2OImageCaptionLoader
-
-        caption_loader = H2OImageCaptionLoader(
-            None,
-            blip_model=captions_model,
-            blip_processor=captions_model,
-            caption_gpu=caption_gpu,
-        ).load_model()
-    else:
-        if enable_captions:
-            caption_loader = "gpu" if caption_gpu else "cpu"
-        else:
-            caption_loader = False
-
-    if verbose:
-        print("Getting sources", flush=True)
-    assert (
-        user_path is not None or url is not None
-    ), "Can't have both user_path and url as None"
-    if not url:
-        assert os.path.isdir(user_path), (
-            "user_path=%s does not exist" % user_path
-        )
-    sources = glob_to_db(
-        user_path,
-        chunk=chunk,
-        chunk_size=chunk_size,
-        verbose=verbose,
-        fail_any_exception=fail_any_exception,
-        n_jobs=n_jobs,
-        url=url,
-        enable_captions=enable_captions,
-        captions_model=captions_model,
-        caption_loader=caption_loader,
-        enable_ocr=enable_ocr,
-    )
-    exceptions = [x for x in sources if x.metadata.get("exception")]
-    print("Exceptions: %s" % exceptions, flush=True)
-    sources = [x for x in sources if "exception" not in x.metadata]
-
-    assert len(sources) > 0, "No sources found"
-    db = create_or_update_db(
-        db_type,
-        persist_directory,
-        collection_name,
-        sources,
-        use_openai_embedding,
-        add_if_exists,
-        verbose,
-        hf_embedding_model,
-    )
-
-    assert db is not None
-    if verbose:
-        print("DONE", flush=True)
-    return db, collection_name

apps/language_models/langchain/read_wiki_full.py

@@ -1,403 +0,0 @@
-"""Load Data from a MediaWiki dump xml."""
-import ast
-import glob
-import pickle
-import uuid
-from typing import List, Optional
-import os
-import bz2
-import csv
-import numpy as np
-import pandas as pd
-import pytest
-from matplotlib import pyplot as plt
-
-from langchain.docstore.document import Document
-from langchain.document_loaders import MWDumpLoader
-
-# path where downloaded wiki files exist, to be processed
-root_path = "/data/jon/h2o-llm"
-
-
-def unescape(x):
-    try:
-        x = ast.literal_eval(x)
-    except:
-        try:
-            x = x.encode("ascii", "ignore").decode("unicode_escape")
-        except:
-            pass
-    return x
-
-
-def get_views():
-    # views = pd.read_csv('wiki_page_views_more_1000month.csv')
-    views = pd.read_csv("wiki_page_views_more_5000month.csv")
-    views.index = views["title"]
-    views = views["views"]
-    views = views.to_dict()
-    views = {str(unescape(str(k))): v for k, v in views.items()}
-    views2 = {k.replace("_", " "): v for k, v in views.items()}
-    # views has _ but pages has " "
-    views.update(views2)
-    return views
-
-
-class MWDumpDirectLoader(MWDumpLoader):
-    def __init__(
-        self,
-        data: str,
-        encoding: Optional[str] = "utf8",
-        title_words_limit=None,
-        use_views=True,
-        verbose=True,
-    ):
-        """Initialize with file path."""
-        self.data = data
-        self.encoding = encoding
-        self.title_words_limit = title_words_limit
-        self.verbose = verbose
-        if use_views:
-            # self.views = get_views()
-            # faster to use global shared values
-            self.views = global_views
-        else:
-            self.views = None
-
-    def load(self) -> List[Document]:
-        """Load from file path."""
-        import mwparserfromhell
-        import mwxml
-
-        dump = mwxml.Dump.from_page_xml(self.data)
-
-        docs = []
-
-        for page in dump.pages:
-            if self.views is not None and page.title not in self.views:
-                if self.verbose:
-                    print("Skipped %s low views" % page.title, flush=True)
-                continue
-            for revision in page:
-                if self.title_words_limit is not None:
-                    num_words = len(" ".join(page.title.split("_")).split(" "))
-                    if num_words > self.title_words_limit:
-                        if self.verbose:
-                            print("Skipped %s" % page.title, flush=True)
-                        continue
-                if self.verbose:
-                    if self.views is not None:
-                        print(
-                            "Kept %s views: %s"
-                            % (page.title, self.views[page.title]),
-                            flush=True,
-                        )
-                    else:
-                        print("Kept %s" % page.title, flush=True)
-
-                code = mwparserfromhell.parse(revision.text)
-                text = code.strip_code(
-                    normalize=True, collapse=True, keep_template_params=False
-                )
-                title_url = str(page.title).replace(" ", "_")
-                metadata = dict(
-                    title=page.title,
-                    source="https://en.wikipedia.org/wiki/" + title_url,
-                    id=page.id,
-                    redirect=page.redirect,
-                    views=self.views[page.title]
-                    if self.views is not None
-                    else -1,
-                )
-                metadata = {k: v for k, v in metadata.items() if v is not None}
-                docs.append(Document(page_content=text, metadata=metadata))
-
-        return docs
-
-
-def search_index(search_term, index_filename):
-    byte_flag = False
-    data_length = start_byte = 0
-    index_file = open(index_filename, "r")
-    csv_reader = csv.reader(index_file, delimiter=":")
-    for line in csv_reader:
-        if not byte_flag and search_term == line[2]:
-            start_byte = int(line[0])
-            byte_flag = True
-        elif byte_flag and int(line[0]) != start_byte:
-            data_length = int(line[0]) - start_byte
-            break
-    index_file.close()
-    return start_byte, data_length
-
-
-def get_start_bytes(index_filename):
-    index_file = open(index_filename, "r")
-    csv_reader = csv.reader(index_file, delimiter=":")
-    start_bytes = set()
-    for line in csv_reader:
-        start_bytes.add(int(line[0]))
-    index_file.close()
-    return sorted(start_bytes)
-
-
-def get_wiki_filenames():
-    # requires
-    # wget http://ftp.acc.umu.se/mirror/wikimedia.org/dumps/enwiki/20230401/enwiki-20230401-pages-articles-multistream-index.txt.bz2
-    base_path = os.path.join(
-        root_path, "enwiki-20230401-pages-articles-multistream"
-    )
-    index_file = "enwiki-20230401-pages-articles-multistream-index.txt"
-    index_filename = os.path.join(base_path, index_file)
-    wiki_filename = os.path.join(
-        base_path, "enwiki-20230401-pages-articles-multistream.xml.bz2"
-    )
-    return index_filename, wiki_filename
-
-
-def get_documents_by_search_term(search_term):
-    index_filename, wiki_filename = get_wiki_filenames()
-    start_byte, data_length = search_index(search_term, index_filename)
-    with open(wiki_filename, "rb") as wiki_file:
-        wiki_file.seek(start_byte)
-        data = bz2.BZ2Decompressor().decompress(wiki_file.read(data_length))
-
-    loader = MWDumpDirectLoader(data.decode())
-    documents = loader.load()
-    return documents
-
-
-def get_one_chunk(
-    wiki_filename,
-    start_byte,
-    end_byte,
-    return_file=True,
-    title_words_limit=None,
-    use_views=True,
-):
-    data_length = end_byte - start_byte
-    with open(wiki_filename, "rb") as wiki_file:
-        wiki_file.seek(start_byte)
-        data = bz2.BZ2Decompressor().decompress(wiki_file.read(data_length))
-
-    loader = MWDumpDirectLoader(
-        data.decode(), title_words_limit=title_words_limit, use_views=use_views
-    )
-    documents1 = loader.load()
-    if return_file:
-        base_tmp = "temp_wiki"
-        if not os.path.isdir(base_tmp):
-            os.makedirs(base_tmp, exist_ok=True)
-        filename = os.path.join(base_tmp, str(uuid.uuid4()) + ".tmp.pickle")
-        with open(filename, "wb") as f:
-            pickle.dump(documents1, f)
-        return filename
-    return documents1
-
-
-from joblib import Parallel, delayed
-
-global_views = get_views()
-
-
-def get_all_documents(small_test=2, n_jobs=None, use_views=True):
-    print("DO get all wiki docs: %s" % small_test, flush=True)
-    index_filename, wiki_filename = get_wiki_filenames()
-    start_bytes = get_start_bytes(index_filename)
-    end_bytes = start_bytes[1:]
-    start_bytes = start_bytes[:-1]
-
-    if small_test:
-        start_bytes = start_bytes[:small_test]
-        end_bytes = end_bytes[:small_test]
-        if n_jobs is None:
-            n_jobs = 5
-    else:
-        if n_jobs is None:
-            n_jobs = os.cpu_count() // 4
-
-    # default loky backend leads to name space conflict problems
-    return_file = True  # large return from joblib hangs
-    documents = Parallel(n_jobs=n_jobs, verbose=10, backend="multiprocessing")(
-        delayed(get_one_chunk)(
-            wiki_filename,
-            start_byte,
-            end_byte,
-            return_file=return_file,
-            use_views=use_views,
-        )
-        for start_byte, end_byte in zip(start_bytes, end_bytes)
-    )
-    if return_file:
-        # then documents really are files
-        files = documents.copy()
-        documents = []
-        for fil in files:
-            with open(fil, "rb") as f:
-                documents.extend(pickle.load(f))
-            os.remove(fil)
-    else:
-        from functools import reduce
-        from operator import concat
-
-        documents = reduce(concat, documents)
-    assert isinstance(documents, list)
-
-    print("DONE get all wiki docs", flush=True)
-    return documents
-
-
-def test_by_search_term():
-    search_term = "Apollo"
-    assert len(get_documents_by_search_term(search_term)) == 100
-
-    search_term = "Abstract (law)"
-    assert len(get_documents_by_search_term(search_term)) == 100
-
-    search_term = "Artificial languages"
-    assert len(get_documents_by_search_term(search_term)) == 100
-
-
-def test_start_bytes():
-    index_filename, wiki_filename = get_wiki_filenames()
-    assert len(get_start_bytes(index_filename)) == 227850
-
-
-def test_get_all_documents():
-    small_test = 20  # 227850
-    n_jobs = os.cpu_count() // 4
-
-    assert (
-        len(
-            get_all_documents(
-                small_test=small_test, n_jobs=n_jobs, use_views=False
-            )
-        )
-        == small_test * 100
-    )
-
-    assert (
-        len(
-            get_all_documents(
-                small_test=small_test, n_jobs=n_jobs, use_views=True
-            )
-        )
-        == 429
-    )
-
-
-def get_one_pageviews(fil):
-    df1 = pd.read_csv(
-        fil,
-        sep=" ",
-        header=None,
-        names=["region", "title", "views", "foo"],
-        quoting=csv.QUOTE_NONE,
-    )
-    df1.index = df1["title"]
-    df1 = df1[df1["region"] == "en"]
-    df1 = df1.drop("region", axis=1)
-    df1 = df1.drop("foo", axis=1)
-    df1 = df1.drop("title", axis=1)  # already index
-
-    base_tmp = "temp_wiki_pageviews"
-    if not os.path.isdir(base_tmp):
-        os.makedirs(base_tmp, exist_ok=True)
-    filename = os.path.join(base_tmp, str(uuid.uuid4()) + ".tmp.csv")
-    df1.to_csv(filename, index=True)
-    return filename
-
-
-def test_agg_pageviews(gen_files=False):
-    if gen_files:
-        path = os.path.join(
-            root_path,
-            "wiki_pageviews/dumps.wikimedia.org/other/pageviews/2023/2023-04",
-        )
-        files = glob.glob(os.path.join(path, "pageviews*.gz"))
-        # files = files[:2]  # test
-        n_jobs = os.cpu_count() // 2
-        csv_files = Parallel(
-            n_jobs=n_jobs, verbose=10, backend="multiprocessing"
-        )(delayed(get_one_pageviews)(fil) for fil in files)
-    else:
-        # to continue without redoing above
-        csv_files = glob.glob(
-            os.path.join(root_path, "temp_wiki_pageviews/*.csv")
-        )
-
-    df_list = []
-    for csv_file in csv_files:
-        print(csv_file)
-        df1 = pd.read_csv(csv_file)
-        df_list.append(df1)
-    df = pd.concat(df_list, axis=0)
-    df = df.groupby("title")["views"].sum().reset_index()
-    df.to_csv("wiki_page_views.csv", index=True)
-
-
-def test_reduce_pageview():
-    filename = "wiki_page_views.csv"
-    df = pd.read_csv(filename)
-    df = df[df["views"] < 1e7]
-    #
-    plt.hist(df["views"], bins=100, log=True)
-    views_avg = np.mean(df["views"])
-    views_median = np.median(df["views"])
-    plt.title("Views avg: %s median: %s" % (views_avg, views_median))
-    plt.savefig(filename.replace(".csv", ".png"))
-    plt.close()
-    #
-    views_limit = 5000
-    df = df[df["views"] > views_limit]
-    filename = "wiki_page_views_more_5000month.csv"
-    df.to_csv(filename, index=True)
-    #
-    plt.hist(df["views"], bins=100, log=True)
-    views_avg = np.mean(df["views"])
-    views_median = np.median(df["views"])
-    plt.title("Views avg: %s median: %s" % (views_avg, views_median))
-    plt.savefig(filename.replace(".csv", ".png"))
-    plt.close()
-
-
-@pytest.mark.skip("Only if doing full processing again, some manual steps")
-def test_do_wiki_full_all():
-    # Install other requirements for wiki specific conversion:
-    # pip install -r reqs_optional/requirements_optional_wikiprocessing.txt
-
-    # Use "Transmission" in Ubuntu to get wiki dump using torrent:
-    # See: https://meta.wikimedia.org/wiki/Data_dump_torrents
-    # E.g. magnet:?xt=urn:btih:b2c74af2b1531d0b63f1166d2011116f44a8fed0&dn=enwiki-20230401-pages-articles-multistream.xml.bz2&tr=udp%3A%2F%2Ftracker.opentrackr.org%3A1337
-
-    # Get index
-    os.system(
-        "wget http://ftp.acc.umu.se/mirror/wikimedia.org/dumps/enwiki/20230401/enwiki-20230401-pages-articles-multistream-index.txt.bz2"
-    )
-
-    # Test that can use LangChain to get docs from subset of wiki as sampled out of full wiki directly using bzip multistream
-    test_get_all_documents()
-
-    # Check can search wiki multistream
-    test_by_search_term()
-
-    # Test can get all start bytes in index
-    test_start_bytes()
-
-    # Get page views, e.g. for entire month of April 2023
-    os.system(
-        "wget -b -m -k -o wget.log -e robots=off https://dumps.wikimedia.org/other/pageviews/2023/2023-04/"
-    )
-
-    # Aggregate page views from many files into single file
-    test_agg_pageviews(gen_files=True)
-
-    # Reduce page views to some limit, so processing of full wiki is not too large
-    test_reduce_pageview()
-
-    # Start generate.py with requesting wiki_full in prep.  This will use page views as referenced in get_views.
-    # Note get_views as global() function done once is required to avoid very slow processing
-    # WARNING: Requires alot of memory to handle, used up to 300GB system RAM at peak
-    """
-    python generate.py --langchain_mode='wiki_full' --visible_langchain_modes="['wiki_full', 'UserData', 'MyData', 'github h2oGPT', 'DriverlessAI docs']" &> lc_out.log
-    """

apps/language_models/langchain/stopping.py

@@ -1,121 +0,0 @@
-import torch
-from transformers import StoppingCriteria, StoppingCriteriaList
-
-from enums import PromptType
-
-
-class StoppingCriteriaSub(StoppingCriteria):
-    def __init__(
-        self, stops=[], encounters=[], device="cuda", model_max_length=None
-    ):
-        super().__init__()
-        assert (
-            len(stops) % len(encounters) == 0
-        ), "Number of stops and encounters must match"
-        self.encounters = encounters
-        self.stops = [stop.to(device) for stop in stops]
-        self.num_stops = [0] * len(stops)
-        self.model_max_length = model_max_length
-
-    def __call__(
-        self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
-    ) -> bool:
-        for stopi, stop in enumerate(self.stops):
-            if torch.all((stop == input_ids[0][-len(stop) :])).item():
-                self.num_stops[stopi] += 1
-                if (
-                    self.num_stops[stopi]
-                    >= self.encounters[stopi % len(self.encounters)]
-                ):
-                    # print("Stopped", flush=True)
-                    return True
-        if (
-            self.model_max_length is not None
-            and input_ids[0].shape[0] >= self.model_max_length
-        ):
-            # critical limit
-            return True
-        # print("Tokens: %s" % input_ids[0].cpu().numpy(), flush=True)
-        # print("Stop Tokens: %s" % [x.cpu().numpy() for x in self.stops], flush=True)
-        return False
-
-
-def get_stopping(
-    prompt_type,
-    prompt_dict,
-    tokenizer,
-    device,
-    human="<human>:",
-    bot="<bot>:",
-    model_max_length=None,
-):
-    # FIXME: prompt_dict unused currently
-    if prompt_type in [
-        PromptType.human_bot.name,
-        PromptType.instruct_vicuna.name,
-        PromptType.instruct_with_end.name,
-    ]:
-        if prompt_type == PromptType.human_bot.name:
-            # encounters = [prompt.count(human) + 1, prompt.count(bot) + 1]
-            # stopping only starts once output is beyond prompt
-            # 1 human is enough to trigger, but need 2 bots, because very first view back will be bot we added
-            stop_words = [human, bot, "\n" + human, "\n" + bot]
-            encounters = [1, 2]
-        elif prompt_type == PromptType.instruct_vicuna.name:
-            # even below is not enough, generic strings and many ways to encode
-            stop_words = [
-                "### Human:",
-                """
-### Human:""",
-                """
-### Human:
-""",
-                "### Assistant:",
-                """
-### Assistant:""",
-                """
-### Assistant:
-""",
-            ]
-            encounters = [1, 2]
-        else:
-            # some instruct prompts have this as end, doesn't hurt to stop on it since not common otherwise
-            stop_words = ["### End"]
-            encounters = [1]
-        stop_words_ids = [
-            tokenizer(stop_word, return_tensors="pt")["input_ids"].squeeze()
-            for stop_word in stop_words
-        ]
-        # handle single token case
-        stop_words_ids = [
-            x if len(x.shape) > 0 else torch.tensor([x])
-            for x in stop_words_ids
-        ]
-        stop_words_ids = [x for x in stop_words_ids if x.shape[0] > 0]
-        # avoid padding in front of tokens
-        if (
-            tokenizer._pad_token
-        ):  # use hidden variable to avoid annoying properly logger bug
-            stop_words_ids = [
-                x[1:] if x[0] == tokenizer.pad_token_id and len(x) > 1 else x
-                for x in stop_words_ids
-            ]
-        # handle fake \n added
-        stop_words_ids = [
-            x[1:] if y[0] == "\n" else x
-            for x, y in zip(stop_words_ids, stop_words)
-        ]
-        # build stopper
-        stopping_criteria = StoppingCriteriaList(
-            [
-                StoppingCriteriaSub(
-                    stops=stop_words_ids,
-                    encounters=encounters,
-                    device=device,
-                    model_max_length=model_max_length,
-                )
-            ]
-        )
-    else:
-        stopping_criteria = StoppingCriteriaList()
-    return stopping_criteria

apps/language_models/langchain/utils_langchain.py

@@ -1,69 +0,0 @@
-from typing import Any, Dict, List, Union, Optional
-import time
-import queue
-
-from langchain.callbacks.base import BaseCallbackHandler
-from langchain.schema import LLMResult
-
-
-class StreamingGradioCallbackHandler(BaseCallbackHandler):
-    """
-    Similar to H2OTextIteratorStreamer that is for HF backend, but here LangChain backend
-    """
-
-    def __init__(self, timeout: Optional[float] = None, block=True):
-        super().__init__()
-        self.text_queue = queue.SimpleQueue()
-        self.stop_signal = None
-        self.do_stop = False
-        self.timeout = timeout
-        self.block = block
-
-    def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> None:
-        """Run when LLM starts running. Clean the queue."""
-        while not self.text_queue.empty():
-            try:
-                self.text_queue.get(block=False)
-            except queue.Empty:
-                continue
-
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Run on new LLM token. Only available when streaming is enabled."""
-        self.text_queue.put(token)
-
-    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        """Run when LLM ends running."""
-        self.text_queue.put(self.stop_signal)
-
-    def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when LLM errors."""
-        self.text_queue.put(self.stop_signal)
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        while True:
-            try:
-                value = (
-                    self.stop_signal
-                )  # value looks unused in pycharm, not true
-                if self.do_stop:
-                    print("hit stop", flush=True)
-                    # could raise or break, maybe best to raise and make parent see if any exception in thread
-                    raise StopIteration()
-                    # break
-                value = self.text_queue.get(
-                    block=self.block, timeout=self.timeout
-                )
-                break
-            except queue.Empty:
-                time.sleep(0.01)
-        if value == self.stop_signal:
-            raise StopIteration()
-        else:
-            return value

apps/language_models/scripts/llama_ir_conversion_utils.py

@@ -1,442 +0,0 @@
-from pathlib import Path
-import argparse
-from argparse import RawTextHelpFormatter
-import re, gc
-
-"""
-    This script can be used as a standalone utility to convert IRs to dynamic + combine them.
-    Following are the various ways this script can be used :-
-        a. To convert a single Linalg IR to dynamic IR:
-            --dynamic --first_ir_path=<PATH TO FIRST IR>
-        b. To convert two Linalg IRs to dynamic IR:
-            --dynamic --first_ir_path=<PATH TO SECOND IR> --first_ir_path=<PATH TO SECOND IR>
-        c. To combine two Linalg IRs into one:
-            --combine --first_ir_path=<PATH TO FIRST IR> --second_ir_path=<PATH TO SECOND IR>
-        d. To convert both IRs into dynamic as well as combine the IRs:
-            --dynamic --combine --first_ir_path=<PATH TO FIRST IR> --second_ir_path=<PATH TO SECOND IR>
-
-    NOTE: For dynamic you'll also need to provide the following set of flags:-
-           i. For First Llama : --dynamic_input_size (DEFAULT: 19)
-          ii. For Second Llama: --model_name (DEFAULT: llama2_7b)
-                                --precision (DEFAULT: 'int4')
-          You may use --save_dynamic to also save the dynamic IR in option d above.
-          Else for option a. and b. the dynamic IR(s) will get saved by default.
-"""
-
-
-def combine_mlir_scripts(
-    first_vicuna_mlir,
-    second_vicuna_mlir,
-    output_name,
-    return_ir=True,
-):
-    print(f"[DEBUG] combining first and second mlir")
-    print(f"[DEBUG] output_name = {output_name}")
-    maps1 = []
-    maps2 = []
-    constants = set()
-    f1 = []
-    f2 = []
-
-    print(f"[DEBUG] processing first vicuna mlir")
-    first_vicuna_mlir = first_vicuna_mlir.splitlines()
-    while first_vicuna_mlir:
-        line = first_vicuna_mlir.pop(0)
-        if re.search("#map\d*\s*=", line):
-            maps1.append(line)
-        elif re.search("arith.constant", line):
-            constants.add(line)
-        elif not re.search("module", line):
-            line = re.sub("forward", "first_vicuna_forward", line)
-            f1.append(line)
-    f1 = f1[:-1]
-    del first_vicuna_mlir
-    gc.collect()
-
-    for i, map_line in enumerate(maps1):
-        map_var = map_line.split(" ")[0]
-        map_line = re.sub(f"{map_var}(?!\d)", map_var + "_0", map_line)
-        maps1[i] = map_line
-        f1 = [
-            re.sub(f"{map_var}(?!\d)", map_var + "_0", func_line)
-            for func_line in f1
-        ]
-
-    print(f"[DEBUG] processing second vicuna mlir")
-    second_vicuna_mlir = second_vicuna_mlir.splitlines()
-    while second_vicuna_mlir:
-        line = second_vicuna_mlir.pop(0)
-        if re.search("#map\d*\s*=", line):
-            maps2.append(line)
-        elif "global_seed" in line:
-            continue
-        elif re.search("arith.constant", line):
-            constants.add(line)
-        elif not re.search("module", line):
-            line = re.sub("forward", "second_vicuna_forward", line)
-            f2.append(line)
-    f2 = f2[:-1]
-    del second_vicuna_mlir
-    gc.collect()
-
-    for i, map_line in enumerate(maps2):
-        map_var = map_line.split(" ")[0]
-        map_line = re.sub(f"{map_var}(?!\d)", map_var + "_1", map_line)
-        maps2[i] = map_line
-        f2 = [
-            re.sub(f"{map_var}(?!\d)", map_var + "_1", func_line)
-            for func_line in f2
-        ]
-
-    module_start = 'module attributes {torch.debug_module_name = "_lambda"} {'
-    module_end = "}"
-
-    global_vars = []
-    vnames = []
-    global_var_loading1 = []
-    global_var_loading2 = []
-
-    print(f"[DEBUG] processing constants")
-    counter = 0
-    constants = list(constants)
-    while constants:
-        constant = constants.pop(0)
-        vname, vbody = constant.split("=")
-        vname = re.sub("%", "", vname)
-        vname = vname.strip()
-        vbody = re.sub("arith.constant", "", vbody)
-        vbody = vbody.strip()
-        if len(vbody.split(":")) < 2:
-            print(constant)
-        vdtype = vbody.split(":")[-1].strip()
-        fixed_vdtype = vdtype
-        if "c1_i64" in vname:
-            print(constant)
-            counter += 1
-        if counter == 2:
-            counter = 0
-            print("detected duplicate")
-            continue
-        vnames.append(vname)
-        if "true" not in vname:
-            global_vars.append(
-                f"ml_program.global private @{vname}({vbody}) : {fixed_vdtype}"
-            )
-            global_var_loading1.append(
-                f"\t\t%{vname} = ml_program.global_load_const @{vname} : {fixed_vdtype}"
-            )
-            global_var_loading2.append(
-                f"\t\t%{vname} = ml_program.global_load_const @{vname} : {fixed_vdtype}"
-            )
-        else:
-            global_vars.append(
-                f"ml_program.global private @{vname}({vbody}) : i1"
-            )
-            global_var_loading1.append(
-                f"\t\t%{vname} = ml_program.global_load_const @{vname} : i1"
-            )
-            global_var_loading2.append(
-                f"\t\t%{vname} = ml_program.global_load_const @{vname} : i1"
-            )
-
-    new_f1, new_f2 = [], []
-
-    print(f"[DEBUG] processing f1")
-    for line in f1:
-        if "func.func" in line:
-            new_f1.append(line)
-            for global_var in global_var_loading1:
-                new_f1.append(global_var)
-        else:
-            new_f1.append(line)
-
-    print(f"[DEBUG] processing f2")
-    for line in f2:
-        if "func.func" in line:
-            new_f2.append(line)
-            for global_var in global_var_loading2:
-                if (
-                    "c20_i64 = arith.addi %dim_i64, %c1_i64 : i64"
-                    in global_var
-                ):
-                    print(global_var)
-                new_f2.append(global_var)
-        else:
-            new_f2.append(line)
-
-    f1 = new_f1
-    f2 = new_f2
-
-    del new_f1
-    del new_f2
-    gc.collect()
-
-    print(
-        [
-            "c20_i64 = arith.addi %dim_i64, %c1_i64 : i64" in x
-            for x in [maps1, maps2, global_vars, f1, f2]
-        ]
-    )
-
-    # doing it this way rather than assembling the whole string
-    # to prevent OOM with 64GiB RAM when encoding the file.
-
-    print(f"[DEBUG] Saving mlir to {output_name}")
-    with open(output_name, "w+") as f_:
-        f_.writelines(line + "\n" for line in maps1)
-        f_.writelines(line + "\n" for line in maps2)
-        f_.writelines(line + "\n" for line in [module_start])
-        f_.writelines(line + "\n" for line in global_vars)
-        f_.writelines(line + "\n" for line in f1)
-        f_.writelines(line + "\n" for line in f2)
-        f_.writelines(line + "\n" for line in [module_end])
-
-    del maps1
-    del maps2
-    del module_start
-    del global_vars
-    del f1
-    del f2
-    del module_end
-    gc.collect()
-
-    if return_ir:
-        print(f"[DEBUG] Reading combined mlir back in")
-        with open(output_name, "rb") as f:
-            return f.read()
-
-
-def write_in_dynamic_inputs0(module, dynamic_input_size):
-    print("[DEBUG] writing dynamic inputs to first vicuna")
-    # Current solution for ensuring mlir files support dynamic inputs
-    # TODO: find a more elegant way to implement this
-    new_lines = []
-    module = module.splitlines()
-    while module:
-        line = module.pop(0)
-        line = re.sub(f"{dynamic_input_size}x", "?x", line)
-        if "?x" in line:
-            line = re.sub("tensor.empty\(\)", "tensor.empty(%dim)", line)
-        line = re.sub(f" {dynamic_input_size},", " %dim,", line)
-        if "tensor.empty" in line and "?x?" in line:
-            line = re.sub(
-                "tensor.empty\(%dim\)", "tensor.empty(%dim, %dim)", line
-            )
-        if "arith.cmpi" in line:
-            line = re.sub(f"c{dynamic_input_size}", "dim", line)
-        if "%0 = tensor.empty(%dim) : tensor<?xi64>" in line:
-            new_lines.append("%dim = tensor.dim %arg0, %c1 : tensor<1x?xi64>")
-        if "%dim = tensor.dim %arg0, %c1 : tensor<1x?xi64>" in line:
-            continue
-
-        new_lines.append(line)
-    return "\n".join(new_lines)
-
-
-def write_in_dynamic_inputs1(module, model_name, precision):
-    print("[DEBUG] writing dynamic inputs to second vicuna")
-
-    def remove_constant_dim(line):
-        if "c19_i64" in line:
-            line = re.sub("c19_i64", "dim_i64", line)
-        if "19x" in line:
-            line = re.sub("19x", "?x", line)
-            line = re.sub("tensor.empty\(\)", "tensor.empty(%dim)", line)
-        if "tensor.empty" in line and "?x?" in line:
-            line = re.sub(
-                "tensor.empty\(%dim\)",
-                "tensor.empty(%dim, %dim)",
-                line,
-            )
-        if "arith.cmpi" in line:
-            line = re.sub("c19", "dim", line)
-        if " 19," in line:
-            line = re.sub(" 19,", " %dim,", line)
-        if "x20x" in line or "<20x" in line:
-            line = re.sub("20x", "?x", line)
-            line = re.sub("tensor.empty\(\)", "tensor.empty(%dimp1)", line)
-        if " 20," in line:
-            line = re.sub(" 20,", " %dimp1,", line)
-        return line
-
-    module = module.splitlines()
-    new_lines = []
-
-    # Using a while loop and the pop method to avoid creating a copy of module
-    if "llama2_13b" in model_name:
-        pkv_tensor_shape = "tensor<1x40x?x128x"
-    elif "llama2_70b" in model_name:
-        pkv_tensor_shape = "tensor<1x8x?x128x"
-    else:
-        pkv_tensor_shape = "tensor<1x32x?x128x"
-    if precision in ["fp16", "int4", "int8"]:
-        pkv_tensor_shape += "f16>"
-    else:
-        pkv_tensor_shape += "f32>"
-
-    while module:
-        line = module.pop(0)
-        if "%c19_i64 = arith.constant 19 : i64" in line:
-            new_lines.append("%c2 = arith.constant 2 : index")
-            new_lines.append(
-                f"%dim_4_int = tensor.dim %arg1, %c2 : {pkv_tensor_shape}"
-            )
-            new_lines.append(
-                "%dim_i64 = arith.index_cast %dim_4_int : index to i64"
-            )
-            continue
-        if "%c2 = arith.constant 2 : index" in line:
-            continue
-        if "%c20_i64 = arith.constant 20 : i64" in line:
-            new_lines.append("%c1_i64 = arith.constant 1 : i64")
-            new_lines.append("%c20_i64 = arith.addi %dim_i64, %c1_i64 : i64")
-            new_lines.append(
-                "%dimp1 = arith.index_cast %c20_i64 : i64 to index"
-            )
-            continue
-        line = remove_constant_dim(line)
-        new_lines.append(line)
-
-    return "\n".join(new_lines)
-
-
-def save_dynamic_ir(ir_to_save, output_file):
-    if not ir_to_save:
-        return
-    # We only get string output from the dynamic conversion utility.
-    from contextlib import redirect_stdout
-
-    with open(output_file, "w") as f:
-        with redirect_stdout(f):
-            print(ir_to_save)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        prog="llama ir utility",
-        description="\tThis script can be used as a standalone utility to convert IRs to dynamic + combine them.\n"
-        + "\tFollowing are the various ways this script can be used :-\n"
-        + "\t\ta. To convert a single Linalg IR to dynamic IR:\n"
-        + "\t\t\t--dynamic --first_ir_path=<PATH TO FIRST IR>\n"
-        + "\t\tb. To convert two Linalg IRs to dynamic IR:\n"
-        + "\t\t\t--dynamic --first_ir_path=<PATH TO SECOND IR> --first_ir_path=<PATH TO SECOND IR>\n"
-        + "\t\tc. To combine two Linalg IRs into one:\n"
-        + "\t\t\t--combine --first_ir_path=<PATH TO FIRST IR> --second_ir_path=<PATH TO SECOND IR>\n"
-        + "\t\td. To convert both IRs into dynamic as well as combine the IRs:\n"
-        + "\t\t\t--dynamic --combine --first_ir_path=<PATH TO FIRST IR> --second_ir_path=<PATH TO SECOND IR>\n\n"
-        + "\tNOTE: For dynamic you'll also need to provide the following set of flags:-\n"
-        + "\t\t i. For First Llama : --dynamic_input_size (DEFAULT: 19)\n"
-        + "\t\tii. For Second Llama: --model_name (DEFAULT: llama2_7b)\n"
-        + "\t\t\t--precision (DEFAULT: 'int4')\n"
-        + "\t      You may use --save_dynamic to also save the dynamic IR in option d above.\n"
-        + "\t      Else for option a. and b. the dynamic IR(s) will get saved by default.\n",
-        formatter_class=RawTextHelpFormatter,
-    )
-    parser.add_argument(
-        "--precision",
-        "-p",
-        default="int4",
-        choices=["fp32", "fp16", "int8", "int4"],
-        help="Precision of the concerned IR",
-    )
-    parser.add_argument(
-        "--model_name",
-        type=str,
-        default="llama2_7b",
-        choices=["vicuna", "llama2_7b", "llama2_13b", "llama2_70b"],
-        help="Specify which model to run.",
-    )
-    parser.add_argument(
-        "--first_ir_path",
-        default=None,
-        help="path to first llama mlir file",
-    )
-    parser.add_argument(
-        "--second_ir_path",
-        default=None,
-        help="path to second llama mlir file",
-    )
-    parser.add_argument(
-        "--dynamic_input_size",
-        type=int,
-        default=19,
-        help="Specify the static input size to replace with dynamic dim.",
-    )
-    parser.add_argument(
-        "--dynamic",
-        default=False,
-        action=argparse.BooleanOptionalAction,
-        help="Converts the IR(s) to dynamic",
-    )
-    parser.add_argument(
-        "--save_dynamic",
-        default=False,
-        action=argparse.BooleanOptionalAction,
-        help="Save the individual IR(s) after converting to dynamic",
-    )
-    parser.add_argument(
-        "--combine",
-        default=False,
-        action=argparse.BooleanOptionalAction,
-        help="Converts the IR(s) to dynamic",
-    )
-
-    args, unknown = parser.parse_known_args()
-
-    dynamic = args.dynamic
-    combine = args.combine
-    assert (
-        dynamic or combine
-    ), "neither `dynamic` nor `combine` flag is turned on"
-    first_ir_path = args.first_ir_path
-    second_ir_path = args.second_ir_path
-    assert first_ir_path or second_ir_path, "no input ir has been provided"
-    if combine:
-        assert (
-            first_ir_path and second_ir_path
-        ), "you will need to provide both IRs to combine"
-    precision = args.precision
-    model_name = args.model_name
-    dynamic_input_size = args.dynamic_input_size
-    save_dynamic = args.save_dynamic
-
-    print(f"Dynamic conversion utility is turned {'ON' if dynamic else 'OFF'}")
-    print(f"Combining IR utility is turned {'ON' if combine else 'OFF'}")
-
-    if dynamic and not combine:
-        save_dynamic = True
-
-    first_ir = None
-    first_dynamic_ir_name = None
-    second_ir = None
-    second_dynamic_ir_name = None
-    if first_ir_path:
-        first_dynamic_ir_name = f"{Path(first_ir_path).stem}_dynamic"
-        with open(first_ir_path, "r") as f:
-            first_ir = f.read()
-    if second_ir_path:
-        second_dynamic_ir_name = f"{Path(second_ir_path).stem}_dynamic"
-        with open(second_ir_path, "r") as f:
-            second_ir = f.read()
-    if dynamic:
-        first_ir = (
-            write_in_dynamic_inputs0(first_ir, dynamic_input_size)
-            if first_ir
-            else None
-        )
-        second_ir = (
-            write_in_dynamic_inputs1(second_ir, model_name, precision)
-            if second_ir
-            else None
-        )
-        if save_dynamic:
-            save_dynamic_ir(first_ir, f"{first_dynamic_ir_name}.mlir")
-            save_dynamic_ir(second_ir, f"{second_dynamic_ir_name}.mlir")
-
-    if combine:
-        combine_mlir_scripts(
-            first_ir,
-            second_ir,
-            f"{model_name}_{precision}.mlir",
-            return_ir=False,
-        )

apps/language_models/scripts/stablelm.py

@@ -46,7 +46,6 @@ def compile_stableLM(
     model_vmfb_name,
     device="cuda",
     precision="fp32",
-    debug=False,
 ):
     from shark.shark_inference import SharkInference
 
@@ -93,7 +92,7 @@ def compile_stableLM(
     shark_module.compile()
 
     path = shark_module.save_module(
-        vmfb_path.parent.absolute(), vmfb_path.stem, debug=debug
+        vmfb_path.parent.absolute(), vmfb_path.stem
     )
     print("Saved vmfb at ", str(path))
 

apps/language_models/shark_llama_cli.spec

@@ -1,94 +0,0 @@
-# -*- mode: python ; coding: utf-8 -*-
-from PyInstaller.utils.hooks import collect_data_files
-from PyInstaller.utils.hooks import collect_submodules
-from PyInstaller.utils.hooks import copy_metadata
-
-import sys ; sys.setrecursionlimit(sys.getrecursionlimit() * 5)
-
-datas = []
-datas += collect_data_files('torch')
-datas += copy_metadata('torch')
-datas += copy_metadata('tqdm')
-datas += copy_metadata('regex')
-datas += copy_metadata('requests')
-datas += copy_metadata('packaging')
-datas += copy_metadata('filelock')
-datas += copy_metadata('numpy')
-datas += copy_metadata('tokenizers')
-datas += copy_metadata('importlib_metadata')
-datas += copy_metadata('torch-mlir')
-datas += copy_metadata('omegaconf')
-datas += copy_metadata('safetensors')
-datas += copy_metadata('huggingface-hub')
-datas += copy_metadata('sentencepiece')
-datas += copy_metadata("pyyaml")
-datas += collect_data_files("tokenizers")
-datas += collect_data_files("tiktoken")
-datas += collect_data_files("accelerate")
-datas += collect_data_files('diffusers')
-datas += collect_data_files('transformers')
-datas += collect_data_files('opencv-python')
-datas += collect_data_files('pytorch_lightning')
-datas += collect_data_files('skimage')
-datas += collect_data_files('gradio')
-datas += collect_data_files('gradio_client')
-datas += collect_data_files('iree')
-datas += collect_data_files('google-cloud-storage')
-datas += collect_data_files('py-cpuinfo')
-datas += collect_data_files("shark", include_py_files=True)
-datas += collect_data_files("timm", include_py_files=True)
-datas += collect_data_files("tqdm")
-datas += collect_data_files("tkinter")
-datas += collect_data_files("webview")
-datas += collect_data_files("sentencepiece")
-datas += collect_data_files("jsonschema")
-datas += collect_data_files("jsonschema_specifications")
-datas += collect_data_files("cpuinfo")
-datas += collect_data_files("langchain")
-
-binaries = []
-
-block_cipher = None
-
-hiddenimports = ['shark', 'shark.shark_inference', 'apps']
-hiddenimports += [x for x in collect_submodules("skimage") if "tests" not in x]
-hiddenimports += [x for x in collect_submodules("iree") if "tests" not in x]
-
-a = Analysis(
-    ['scripts/vicuna.py'],
-    pathex=['.'],
-    binaries=binaries,
-    datas=datas,
-    hiddenimports=hiddenimports,
-    hookspath=[],
-    hooksconfig={},
-    runtime_hooks=[],
-    excludes=[],
-    win_no_prefer_redirects=False,
-    win_private_assemblies=False,
-    cipher=block_cipher,
-    noarchive=False,
-)
-pyz = PYZ(a.pure, a.zipped_data, cipher=block_cipher)
-
-exe = EXE(
-    pyz,
-    a.scripts,
-    a.binaries,
-    a.zipfiles,
-    a.datas,
-    [],
-    name='shark_llama_cli',
-    debug=False,
-    bootloader_ignore_signals=False,
-    strip=False,
-    upx=True,
-    upx_exclude=[],
-    runtime_tmpdir=None,
-    console=True,
-    disable_windowed_traceback=False,
-    argv_emulation=False,
-    target_arch=None,
-    codesign_identity=None,
-    entitlements_file=None,
-)

apps/language_models/src/model_wrappers/falcon_sharded_model.py

@@ -1,675 +0,0 @@
-import torch
-from typing import Optional, Tuple
-
-
-class WordEmbeddingsLayer(torch.nn.Module):
-    def __init__(self, word_embedding_layer):
-        super().__init__()
-        self.model = word_embedding_layer
-
-    def forward(self, input_ids):
-        output = self.model.forward(input=input_ids)
-        return output
-
-
-class CompiledWordEmbeddingsLayer(torch.nn.Module):
-    def __init__(self, compiled_word_embedding_layer):
-        super().__init__()
-        self.model = compiled_word_embedding_layer
-
-    def forward(self, input_ids):
-        input_ids = input_ids.detach().numpy()
-        new_input_ids = self.model("forward", input_ids)
-        new_input_ids = new_input_ids.reshape(
-            [1, new_input_ids.shape[0], new_input_ids.shape[1]]
-        )
-        return torch.tensor(new_input_ids)
-
-
-class LNFEmbeddingLayer(torch.nn.Module):
-    def __init__(self, ln_f):
-        super().__init__()
-        self.model = ln_f
-
-    def forward(self, hidden_states):
-        output = self.model.forward(input=hidden_states)
-        return output
-
-
-class CompiledLNFEmbeddingLayer(torch.nn.Module):
-    def __init__(self, ln_f):
-        super().__init__()
-        self.model = ln_f
-
-    def forward(self, hidden_states):
-        hidden_states = hidden_states.detach().numpy()
-        new_hidden_states = self.model("forward", (hidden_states,))
-
-        return torch.tensor(new_hidden_states)
-
-
-class LMHeadEmbeddingLayer(torch.nn.Module):
-    def __init__(self, embedding_layer):
-        super().__init__()
-        self.model = embedding_layer
-
-    def forward(self, hidden_states):
-        output = self.model.forward(input=hidden_states)
-        return output
-
-
-class CompiledLMHeadEmbeddingLayer(torch.nn.Module):
-    def __init__(self, lm_head):
-        super().__init__()
-        self.model = lm_head
-
-    def forward(self, hidden_states):
-        hidden_states = hidden_states.detach().numpy()
-        new_hidden_states = self.model("forward", (hidden_states,))
-        return torch.tensor(new_hidden_states)
-
-
-class FourWayShardingDecoderLayer(torch.nn.Module):
-    def __init__(self, decoder_layer_model, falcon_variant):
-        super().__init__()
-        self.model = decoder_layer_model
-        self.falcon_variant = falcon_variant
-
-    def forward(self, hidden_states, attention_mask):
-        new_pkvs = []
-        for layer in self.model:
-            outputs = layer(
-                hidden_states=hidden_states,
-                alibi=None,
-                attention_mask=attention_mask,
-                use_cache=True,
-            )
-            hidden_states = outputs[0]
-            new_pkvs.append(
-                (
-                    outputs[-1][0],
-                    outputs[-1][1],
-                )
-            )
-
-        (
-            (new_pkv00, new_pkv01),
-            (new_pkv10, new_pkv11),
-            (new_pkv20, new_pkv21),
-            (new_pkv30, new_pkv31),
-            (new_pkv40, new_pkv41),
-            (new_pkv50, new_pkv51),
-            (new_pkv60, new_pkv61),
-            (new_pkv70, new_pkv71),
-            (new_pkv80, new_pkv81),
-            (new_pkv90, new_pkv91),
-            (new_pkv100, new_pkv101),
-            (new_pkv110, new_pkv111),
-            (new_pkv120, new_pkv121),
-            (new_pkv130, new_pkv131),
-            (new_pkv140, new_pkv141),
-            (new_pkv150, new_pkv151),
-            (new_pkv160, new_pkv161),
-            (new_pkv170, new_pkv171),
-            (new_pkv180, new_pkv181),
-            (new_pkv190, new_pkv191),
-        ) = new_pkvs
-        result = (
-            hidden_states,
-            new_pkv00,
-            new_pkv01,
-            new_pkv10,
-            new_pkv11,
-            new_pkv20,
-            new_pkv21,
-            new_pkv30,
-            new_pkv31,
-            new_pkv40,
-            new_pkv41,
-            new_pkv50,
-            new_pkv51,
-            new_pkv60,
-            new_pkv61,
-            new_pkv70,
-            new_pkv71,
-            new_pkv80,
-            new_pkv81,
-            new_pkv90,
-            new_pkv91,
-            new_pkv100,
-            new_pkv101,
-            new_pkv110,
-            new_pkv111,
-            new_pkv120,
-            new_pkv121,
-            new_pkv130,
-            new_pkv131,
-            new_pkv140,
-            new_pkv141,
-            new_pkv150,
-            new_pkv151,
-            new_pkv160,
-            new_pkv161,
-            new_pkv170,
-            new_pkv171,
-            new_pkv180,
-            new_pkv181,
-            new_pkv190,
-            new_pkv191,
-        )
-        return result
-
-
-class CompiledFourWayShardingDecoderLayer(torch.nn.Module):
-    def __init__(
-        self, layer_id, device_idx, falcon_variant, device, precision, model
-    ):
-        super().__init__()
-        self.layer_id = layer_id
-        self.device_index = device_idx
-        self.falcon_variant = falcon_variant
-        self.device = device
-        self.precision = precision
-        self.model = model
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        alibi: Optional[torch.Tensor],
-        attention_mask: torch.Tensor,
-        position_ids: Optional[torch.LongTensor] = None,
-        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        use_cache: bool = False,
-        output_attentions: bool = False,
-    ):
-        import gc
-
-        torch.cuda.empty_cache()
-        gc.collect()
-
-        if self.model is None:
-            raise ValueError("Layer vmfb not found")
-
-        hidden_states = hidden_states.to(torch.float32).detach().numpy()
-        attention_mask = attention_mask.to(torch.float32).detach().numpy()
-
-        if alibi is not None or layer_past is not None:
-            raise ValueError("Past Key Values and alibi should be None")
-        else:
-            output = self.model(
-                "forward",
-                (
-                    hidden_states,
-                    attention_mask,
-                ),
-            )
-
-        result = (
-            torch.tensor(output[0]),
-            (
-                torch.tensor(output[1]),
-                torch.tensor(output[2]),
-            ),
-            (
-                torch.tensor(output[3]),
-                torch.tensor(output[4]),
-            ),
-            (
-                torch.tensor(output[5]),
-                torch.tensor(output[6]),
-            ),
-            (
-                torch.tensor(output[7]),
-                torch.tensor(output[8]),
-            ),
-            (
-                torch.tensor(output[9]),
-                torch.tensor(output[10]),
-            ),
-            (
-                torch.tensor(output[11]),
-                torch.tensor(output[12]),
-            ),
-            (
-                torch.tensor(output[13]),
-                torch.tensor(output[14]),
-            ),
-            (
-                torch.tensor(output[15]),
-                torch.tensor(output[16]),
-            ),
-            (
-                torch.tensor(output[17]),
-                torch.tensor(output[18]),
-            ),
-            (
-                torch.tensor(output[19]),
-                torch.tensor(output[20]),
-            ),
-            (
-                torch.tensor(output[21]),
-                torch.tensor(output[22]),
-            ),
-            (
-                torch.tensor(output[23]),
-                torch.tensor(output[24]),
-            ),
-            (
-                torch.tensor(output[25]),
-                torch.tensor(output[26]),
-            ),
-            (
-                torch.tensor(output[27]),
-                torch.tensor(output[28]),
-            ),
-            (
-                torch.tensor(output[29]),
-                torch.tensor(output[30]),
-            ),
-            (
-                torch.tensor(output[31]),
-                torch.tensor(output[32]),
-            ),
-            (
-                torch.tensor(output[33]),
-                torch.tensor(output[34]),
-            ),
-            (
-                torch.tensor(output[35]),
-                torch.tensor(output[36]),
-            ),
-            (
-                torch.tensor(output[37]),
-                torch.tensor(output[38]),
-            ),
-            (
-                torch.tensor(output[39]),
-                torch.tensor(output[40]),
-            ),
-        )
-        return result
-
-
-class TwoWayShardingDecoderLayer(torch.nn.Module):
-    def __init__(self, decoder_layer_model, falcon_variant):
-        super().__init__()
-        self.model = decoder_layer_model
-        self.falcon_variant = falcon_variant
-
-    def forward(self, hidden_states, attention_mask):
-        new_pkvs = []
-        for layer in self.model:
-            outputs = layer(
-                hidden_states=hidden_states,
-                alibi=None,
-                attention_mask=attention_mask,
-                use_cache=True,
-            )
-            hidden_states = outputs[0]
-            new_pkvs.append(
-                (
-                    outputs[-1][0],
-                    outputs[-1][1],
-                )
-            )
-
-        (
-            (new_pkv00, new_pkv01),
-            (new_pkv10, new_pkv11),
-            (new_pkv20, new_pkv21),
-            (new_pkv30, new_pkv31),
-            (new_pkv40, new_pkv41),
-            (new_pkv50, new_pkv51),
-            (new_pkv60, new_pkv61),
-            (new_pkv70, new_pkv71),
-            (new_pkv80, new_pkv81),
-            (new_pkv90, new_pkv91),
-            (new_pkv100, new_pkv101),
-            (new_pkv110, new_pkv111),
-            (new_pkv120, new_pkv121),
-            (new_pkv130, new_pkv131),
-            (new_pkv140, new_pkv141),
-            (new_pkv150, new_pkv151),
-            (new_pkv160, new_pkv161),
-            (new_pkv170, new_pkv171),
-            (new_pkv180, new_pkv181),
-            (new_pkv190, new_pkv191),
-            (new_pkv200, new_pkv201),
-            (new_pkv210, new_pkv211),
-            (new_pkv220, new_pkv221),
-            (new_pkv230, new_pkv231),
-            (new_pkv240, new_pkv241),
-            (new_pkv250, new_pkv251),
-            (new_pkv260, new_pkv261),
-            (new_pkv270, new_pkv271),
-            (new_pkv280, new_pkv281),
-            (new_pkv290, new_pkv291),
-            (new_pkv300, new_pkv301),
-            (new_pkv310, new_pkv311),
-            (new_pkv320, new_pkv321),
-            (new_pkv330, new_pkv331),
-            (new_pkv340, new_pkv341),
-            (new_pkv350, new_pkv351),
-            (new_pkv360, new_pkv361),
-            (new_pkv370, new_pkv371),
-            (new_pkv380, new_pkv381),
-            (new_pkv390, new_pkv391),
-        ) = new_pkvs
-        result = (
-            hidden_states,
-            new_pkv00,
-            new_pkv01,
-            new_pkv10,
-            new_pkv11,
-            new_pkv20,
-            new_pkv21,
-            new_pkv30,
-            new_pkv31,
-            new_pkv40,
-            new_pkv41,
-            new_pkv50,
-            new_pkv51,
-            new_pkv60,
-            new_pkv61,
-            new_pkv70,
-            new_pkv71,
-            new_pkv80,
-            new_pkv81,
-            new_pkv90,
-            new_pkv91,
-            new_pkv100,
-            new_pkv101,
-            new_pkv110,
-            new_pkv111,
-            new_pkv120,
-            new_pkv121,
-            new_pkv130,
-            new_pkv131,
-            new_pkv140,
-            new_pkv141,
-            new_pkv150,
-            new_pkv151,
-            new_pkv160,
-            new_pkv161,
-            new_pkv170,
-            new_pkv171,
-            new_pkv180,
-            new_pkv181,
-            new_pkv190,
-            new_pkv191,
-            new_pkv200,
-            new_pkv201,
-            new_pkv210,
-            new_pkv211,
-            new_pkv220,
-            new_pkv221,
-            new_pkv230,
-            new_pkv231,
-            new_pkv240,
-            new_pkv241,
-            new_pkv250,
-            new_pkv251,
-            new_pkv260,
-            new_pkv261,
-            new_pkv270,
-            new_pkv271,
-            new_pkv280,
-            new_pkv281,
-            new_pkv290,
-            new_pkv291,
-            new_pkv300,
-            new_pkv301,
-            new_pkv310,
-            new_pkv311,
-            new_pkv320,
-            new_pkv321,
-            new_pkv330,
-            new_pkv331,
-            new_pkv340,
-            new_pkv341,
-            new_pkv350,
-            new_pkv351,
-            new_pkv360,
-            new_pkv361,
-            new_pkv370,
-            new_pkv371,
-            new_pkv380,
-            new_pkv381,
-            new_pkv390,
-            new_pkv391,
-        )
-        return result
-
-
-class CompiledTwoWayShardingDecoderLayer(torch.nn.Module):
-    def __init__(
-        self, layer_id, device_idx, falcon_variant, device, precision, model
-    ):
-        super().__init__()
-        self.layer_id = layer_id
-        self.device_index = device_idx
-        self.falcon_variant = falcon_variant
-        self.device = device
-        self.precision = precision
-        self.model = model
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        alibi: Optional[torch.Tensor],
-        attention_mask: torch.Tensor,
-        position_ids: Optional[torch.LongTensor] = None,
-        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        use_cache: bool = False,
-        output_attentions: bool = False,
-    ):
-        import gc
-
-        torch.cuda.empty_cache()
-        gc.collect()
-
-        if self.model is None:
-            raise ValueError("Layer vmfb not found")
-
-        hidden_states = hidden_states.to(torch.float32).detach().numpy()
-        attention_mask = attention_mask.to(torch.float32).detach().numpy()
-
-        if alibi is not None or layer_past is not None:
-            raise ValueError("Past Key Values and alibi should be None")
-        else:
-            output = self.model(
-                "forward",
-                (
-                    hidden_states,
-                    attention_mask,
-                ),
-            )
-
-        result = (
-            torch.tensor(output[0]),
-            (
-                torch.tensor(output[1]),
-                torch.tensor(output[2]),
-            ),
-            (
-                torch.tensor(output[3]),
-                torch.tensor(output[4]),
-            ),
-            (
-                torch.tensor(output[5]),
-                torch.tensor(output[6]),
-            ),
-            (
-                torch.tensor(output[7]),
-                torch.tensor(output[8]),
-            ),
-            (
-                torch.tensor(output[9]),
-                torch.tensor(output[10]),
-            ),
-            (
-                torch.tensor(output[11]),
-                torch.tensor(output[12]),
-            ),
-            (
-                torch.tensor(output[13]),
-                torch.tensor(output[14]),
-            ),
-            (
-                torch.tensor(output[15]),
-                torch.tensor(output[16]),
-            ),
-            (
-                torch.tensor(output[17]),
-                torch.tensor(output[18]),
-            ),
-            (
-                torch.tensor(output[19]),
-                torch.tensor(output[20]),
-            ),
-            (
-                torch.tensor(output[21]),
-                torch.tensor(output[22]),
-            ),
-            (
-                torch.tensor(output[23]),
-                torch.tensor(output[24]),
-            ),
-            (
-                torch.tensor(output[25]),
-                torch.tensor(output[26]),
-            ),
-            (
-                torch.tensor(output[27]),
-                torch.tensor(output[28]),
-            ),
-            (
-                torch.tensor(output[29]),
-                torch.tensor(output[30]),
-            ),
-            (
-                torch.tensor(output[31]),
-                torch.tensor(output[32]),
-            ),
-            (
-                torch.tensor(output[33]),
-                torch.tensor(output[34]),
-            ),
-            (
-                torch.tensor(output[35]),
-                torch.tensor(output[36]),
-            ),
-            (
-                torch.tensor(output[37]),
-                torch.tensor(output[38]),
-            ),
-            (
-                torch.tensor(output[39]),
-                torch.tensor(output[40]),
-            ),
-            (
-                torch.tensor(output[41]),
-                torch.tensor(output[42]),
-            ),
-            (
-                torch.tensor(output[43]),
-                torch.tensor(output[44]),
-            ),
-            (
-                torch.tensor(output[45]),
-                torch.tensor(output[46]),
-            ),
-            (
-                torch.tensor(output[47]),
-                torch.tensor(output[48]),
-            ),
-            (
-                torch.tensor(output[49]),
-                torch.tensor(output[50]),
-            ),
-            (
-                torch.tensor(output[51]),
-                torch.tensor(output[52]),
-            ),
-            (
-                torch.tensor(output[53]),
-                torch.tensor(output[54]),
-            ),
-            (
-                torch.tensor(output[55]),
-                torch.tensor(output[56]),
-            ),
-            (
-                torch.tensor(output[57]),
-                torch.tensor(output[58]),
-            ),
-            (
-                torch.tensor(output[59]),
-                torch.tensor(output[60]),
-            ),
-            (
-                torch.tensor(output[61]),
-                torch.tensor(output[62]),
-            ),
-            (
-                torch.tensor(output[63]),
-                torch.tensor(output[64]),
-            ),
-            (
-                torch.tensor(output[65]),
-                torch.tensor(output[66]),
-            ),
-            (
-                torch.tensor(output[67]),
-                torch.tensor(output[68]),
-            ),
-            (
-                torch.tensor(output[69]),
-                torch.tensor(output[70]),
-            ),
-            (
-                torch.tensor(output[71]),
-                torch.tensor(output[72]),
-            ),
-            (
-                torch.tensor(output[73]),
-                torch.tensor(output[74]),
-            ),
-            (
-                torch.tensor(output[75]),
-                torch.tensor(output[76]),
-            ),
-            (
-                torch.tensor(output[77]),
-                torch.tensor(output[78]),
-            ),
-            (
-                torch.tensor(output[79]),
-                torch.tensor(output[80]),
-            ),
-        )
-        return result
-
-
-class ShardedFalconModel:
-    def __init__(self, model, layers, word_embeddings, ln_f, lm_head):
-        super().__init__()
-        self.model = model
-        self.model.transformer.h = torch.nn.modules.container.ModuleList(
-            layers
-        )
-        self.model.transformer.word_embeddings = word_embeddings
-        self.model.transformer.ln_f = ln_f
-        self.model.lm_head = lm_head
-
-    def forward(
-        self,
-        input_ids,
-        attention_mask=None,
-    ):
-        return self.model.forward(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-        ).logits[:, -1, :]

apps/language_models/src/model_wrappers/minigpt4.py

@@ -1,503 +0,0 @@
-import torch
-import dataclasses
-from enum import auto, Enum
-from typing import List, Any
-from transformers import StoppingCriteria
-
-
-from brevitas_examples.common.generative.quantize import quantize_model
-from brevitas_examples.llm.llm_quant.run_utils import get_model_impl
-
-
-class LayerNorm(torch.nn.LayerNorm):
-    """Subclass torch's LayerNorm to handle fp16."""
-
-    def forward(self, x: torch.Tensor):
-        orig_type = x.dtype
-        ret = super().forward(x.type(torch.float32))
-        return ret.type(orig_type)
-
-
-class VisionModel(torch.nn.Module):
-    def __init__(
-        self,
-        ln_vision,
-        visual_encoder,
-        precision="fp32",
-        weight_group_size=128,
-    ):
-        super().__init__()
-        self.ln_vision = ln_vision
-        self.visual_encoder = visual_encoder
-        if precision in ["int4", "int8"]:
-            print("Vision Model applying weight quantization to ln_vision")
-            weight_bit_width = 4 if precision == "int4" else 8
-            quantize_model(
-                self.ln_vision,
-                dtype=torch.float32,
-                weight_bit_width=weight_bit_width,
-                weight_param_method="stats",
-                weight_scale_precision="float_scale",
-                weight_quant_type="asym",
-                weight_quant_granularity="per_group",
-                weight_group_size=weight_group_size,
-                quantize_weight_zero_point=False,
-            )
-            print("Weight quantization applied.")
-            print(
-                "Vision Model applying weight quantization to visual_encoder"
-            )
-            quantize_model(
-                self.visual_encoder,
-                dtype=torch.float32,
-                weight_bit_width=weight_bit_width,
-                weight_param_method="stats",
-                weight_scale_precision="float_scale",
-                weight_quant_type="asym",
-                weight_quant_granularity="per_group",
-                weight_group_size=weight_group_size,
-                quantize_weight_zero_point=False,
-            )
-            print("Weight quantization applied.")
-
-    def forward(self, image):
-        image_embeds = self.ln_vision(self.visual_encoder(image))
-        return image_embeds
-
-
-class QformerBertModel(torch.nn.Module):
-    def __init__(self, qformer_bert):
-        super().__init__()
-        self.qformer_bert = qformer_bert
-
-    def forward(self, query_tokens, image_embeds, image_atts):
-        query_output = self.qformer_bert(
-            query_embeds=query_tokens,
-            encoder_hidden_states=image_embeds,
-            encoder_attention_mask=image_atts,
-            return_dict=True,
-        )
-        return query_output.last_hidden_state
-
-
-class FirstLlamaModel(torch.nn.Module):
-    def __init__(self, model, precision="fp32", weight_group_size=128):
-        super().__init__()
-        self.model = model
-        print("SHARK: Loading LLAMA Done")
-        if precision in ["int4", "int8"]:
-            print("First Llama applying weight quantization")
-            weight_bit_width = 4 if precision == "int4" else 8
-            quantize_model(
-                self.model,
-                dtype=torch.float32,
-                weight_bit_width=weight_bit_width,
-                weight_param_method="stats",
-                weight_scale_precision="float_scale",
-                weight_quant_type="asym",
-                weight_quant_granularity="per_group",
-                weight_group_size=weight_group_size,
-                quantize_weight_zero_point=False,
-            )
-            print("Weight quantization applied.")
-
-    def forward(self, inputs_embeds, position_ids, attention_mask):
-        print("************************************")
-        print(
-            "inputs_embeds: ",
-            inputs_embeds.shape,
-            " dtype: ",
-            inputs_embeds.dtype,
-        )
-        print(
-            "position_ids: ",
-            position_ids.shape,
-            " dtype: ",
-            position_ids.dtype,
-        )
-        print(
-            "attention_mask: ",
-            attention_mask.shape,
-            " dtype: ",
-            attention_mask.dtype,
-        )
-        print("************************************")
-        config = {
-            "inputs_embeds": inputs_embeds,
-            "position_ids": position_ids,
-            "past_key_values": None,
-            "use_cache": True,
-            "attention_mask": attention_mask,
-        }
-        output = self.model(
-            **config,
-            return_dict=True,
-            output_attentions=False,
-            output_hidden_states=False,
-        )
-        return_vals = []
-        return_vals.append(output.logits)
-        temp_past_key_values = output.past_key_values
-        for item in temp_past_key_values:
-            return_vals.append(item[0])
-            return_vals.append(item[1])
-        return tuple(return_vals)
-
-
-class SecondLlamaModel(torch.nn.Module):
-    def __init__(self, model, precision="fp32", weight_group_size=128):
-        super().__init__()
-        self.model = model
-        print("SHARK: Loading LLAMA Done")
-        if precision in ["int4", "int8"]:
-            print("Second Llama applying weight quantization")
-            weight_bit_width = 4 if precision == "int4" else 8
-            quantize_model(
-                self.model,
-                dtype=torch.float32,
-                weight_bit_width=weight_bit_width,
-                weight_param_method="stats",
-                weight_scale_precision="float_scale",
-                weight_quant_type="asym",
-                weight_quant_granularity="per_group",
-                weight_group_size=weight_group_size,
-                quantize_weight_zero_point=False,
-            )
-            print("Weight quantization applied.")
-
-    def forward(
-        self,
-        input_ids,
-        position_ids,
-        attention_mask,
-        i1,
-        i2,
-        i3,
-        i4,
-        i5,
-        i6,
-        i7,
-        i8,
-        i9,
-        i10,
-        i11,
-        i12,
-        i13,
-        i14,
-        i15,
-        i16,
-        i17,
-        i18,
-        i19,
-        i20,
-        i21,
-        i22,
-        i23,
-        i24,
-        i25,
-        i26,
-        i27,
-        i28,
-        i29,
-        i30,
-        i31,
-        i32,
-        i33,
-        i34,
-        i35,
-        i36,
-        i37,
-        i38,
-        i39,
-        i40,
-        i41,
-        i42,
-        i43,
-        i44,
-        i45,
-        i46,
-        i47,
-        i48,
-        i49,
-        i50,
-        i51,
-        i52,
-        i53,
-        i54,
-        i55,
-        i56,
-        i57,
-        i58,
-        i59,
-        i60,
-        i61,
-        i62,
-        i63,
-        i64,
-    ):
-        print("************************************")
-        print("input_ids: ", input_ids.shape, " dtype: ", input_ids.dtype)
-        print(
-            "position_ids: ",
-            position_ids.shape,
-            " dtype: ",
-            position_ids.dtype,
-        )
-        print(
-            "attention_mask: ",
-            attention_mask.shape,
-            " dtype: ",
-            attention_mask.dtype,
-        )
-        print("past_key_values: ", i1.shape, i2.shape, i63.shape, i64.shape)
-        print("past_key_values dtype: ", i1.dtype)
-        print("************************************")
-        config = {
-            "input_ids": input_ids,
-            "position_ids": position_ids,
-            "past_key_values": (
-                (i1, i2),
-                (
-                    i3,
-                    i4,
-                ),
-                (
-                    i5,
-                    i6,
-                ),
-                (
-                    i7,
-                    i8,
-                ),
-                (
-                    i9,
-                    i10,
-                ),
-                (
-                    i11,
-                    i12,
-                ),
-                (
-                    i13,
-                    i14,
-                ),
-                (
-                    i15,
-                    i16,
-                ),
-                (
-                    i17,
-                    i18,
-                ),
-                (
-                    i19,
-                    i20,
-                ),
-                (
-                    i21,
-                    i22,
-                ),
-                (
-                    i23,
-                    i24,
-                ),
-                (
-                    i25,
-                    i26,
-                ),
-                (
-                    i27,
-                    i28,
-                ),
-                (
-                    i29,
-                    i30,
-                ),
-                (
-                    i31,
-                    i32,
-                ),
-                (
-                    i33,
-                    i34,
-                ),
-                (
-                    i35,
-                    i36,
-                ),
-                (
-                    i37,
-                    i38,
-                ),
-                (
-                    i39,
-                    i40,
-                ),
-                (
-                    i41,
-                    i42,
-                ),
-                (
-                    i43,
-                    i44,
-                ),
-                (
-                    i45,
-                    i46,
-                ),
-                (
-                    i47,
-                    i48,
-                ),
-                (
-                    i49,
-                    i50,
-                ),
-                (
-                    i51,
-                    i52,
-                ),
-                (
-                    i53,
-                    i54,
-                ),
-                (
-                    i55,
-                    i56,
-                ),
-                (
-                    i57,
-                    i58,
-                ),
-                (
-                    i59,
-                    i60,
-                ),
-                (
-                    i61,
-                    i62,
-                ),
-                (
-                    i63,
-                    i64,
-                ),
-            ),
-            "use_cache": True,
-            "attention_mask": attention_mask,
-        }
-        output = self.model(
-            **config,
-            return_dict=True,
-            output_attentions=False,
-            output_hidden_states=False,
-        )
-        return_vals = []
-        return_vals.append(output.logits)
-        temp_past_key_values = output.past_key_values
-        for item in temp_past_key_values:
-            return_vals.append(item[0])
-            return_vals.append(item[1])
-        return tuple(return_vals)
-
-
-class SeparatorStyle(Enum):
-    """Different separator style."""
-
-    SINGLE = auto()
-    TWO = auto()
-
-
-@dataclasses.dataclass
-class Conversation:
-    """A class that keeps all conversation history."""
-
-    system: str
-    roles: List[str]
-    messages: List[List[str]]
-    offset: int
-    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
-    sep: str = "###"
-    sep2: str = None
-
-    skip_next: bool = False
-    conv_id: Any = None
-
-    def get_prompt(self):
-        if self.sep_style == SeparatorStyle.SINGLE:
-            ret = self.system + self.sep
-            for role, message in self.messages:
-                if message:
-                    ret += role + ": " + message + self.sep
-                else:
-                    ret += role + ":"
-            return ret
-        elif self.sep_style == SeparatorStyle.TWO:
-            seps = [self.sep, self.sep2]
-            ret = self.system + seps[0]
-            for i, (role, message) in enumerate(self.messages):
-                if message:
-                    ret += role + ": " + message + seps[i % 2]
-                else:
-                    ret += role + ":"
-            return ret
-        else:
-            raise ValueError(f"Invalid style: {self.sep_style}")
-
-    def append_message(self, role, message):
-        self.messages.append([role, message])
-
-    def to_gradio_chatbot(self):
-        ret = []
-        for i, (role, msg) in enumerate(self.messages[self.offset :]):
-            if i % 2 == 0:
-                ret.append([msg, None])
-            else:
-                ret[-1][-1] = msg
-        return ret
-
-    def copy(self):
-        return Conversation(
-            system=self.system,
-            roles=self.roles,
-            messages=[[x, y] for x, y in self.messages],
-            offset=self.offset,
-            sep_style=self.sep_style,
-            sep=self.sep,
-            sep2=self.sep2,
-            conv_id=self.conv_id,
-        )
-
-    def dict(self):
-        return {
-            "system": self.system,
-            "roles": self.roles,
-            "messages": self.messages,
-            "offset": self.offset,
-            "sep": self.sep,
-            "sep2": self.sep2,
-            "conv_id": self.conv_id,
-        }
-
-
-class StoppingCriteriaSub(StoppingCriteria):
-    def __init__(self, stops=[], encounters=1):
-        super().__init__()
-        self.stops = stops
-
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
-        for stop in self.stops:
-            if torch.all((stop == input_ids[0][-len(stop) :])).item():
-                return True
-
-        return False
-
-
-CONV_VISION = Conversation(
-    system="Give the following image: <Img>ImageContent</Img>. "
-    "You will be able to see the image once I provide it to you. Please answer my questions.",
-    roles=("Human", "Assistant"),
-    messages=[],
-    offset=2,
-    sep_style=SeparatorStyle.SINGLE,
-    sep="###",
-)

apps/language_models/src/model_wrappers/vicuna4.py

@@ -1,876 +0,0 @@
-import argparse
-import json
-import re
-from io import BytesIO
-from pathlib import Path
-from tqdm import tqdm
-from typing import List, Optional, Tuple, Union
-import numpy as np
-import iree.runtime
-import itertools
-import subprocess
-
-import torch
-import torch_mlir
-from torch_mlir import TensorPlaceholder
-from torch_mlir.compiler_utils import run_pipeline_with_repro_report
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    LlamaPreTrainedModel,
-)
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-    SequenceClassifierOutputWithPast,
-)
-from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import (
-    add_start_docstrings,
-    add_start_docstrings_to_model_forward,
-    logging,
-    replace_return_docstrings,
-)
-
-from apps.language_models.src.pipelines.SharkLLMBase import SharkLLMBase
-from apps.language_models.src.model_wrappers.vicuna_sharded_model import (
-    FirstVicunaLayer,
-    SecondVicunaLayer,
-    CompiledVicunaLayer,
-    ShardedVicunaModel,
-    LMHead,
-    LMHeadCompiled,
-    VicunaEmbedding,
-    VicunaEmbeddingCompiled,
-    VicunaNorm,
-    VicunaNormCompiled,
-)
-from apps.language_models.src.model_wrappers.vicuna_model import (
-    FirstVicuna,
-    SecondVicuna7B,
-)
-from apps.language_models.utils import (
-    get_vmfb_from_path,
-)
-from shark.shark_downloader import download_public_file
-from shark.shark_importer import get_f16_inputs
-from shark.shark_inference import SharkInference
-
-from transformers.models.llama.configuration_llama import LlamaConfig
-from transformers.models.llama.modeling_llama import (
-    LlamaDecoderLayer,
-    LlamaRMSNorm,
-    _make_causal_mask,
-    _expand_mask,
-)
-from torch import nn
-from time import time
-
-
-class LlamaModel(LlamaPreTrainedModel):
-    """
-    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]
-
-    Args:
-        config: LlamaConfig
-    """
-
-    def __init__(self, config: LlamaConfig):
-        super().__init__(config)
-        self.padding_idx = config.pad_token_id
-        self.vocab_size = config.vocab_size
-
-        self.embed_tokens = nn.Embedding(
-            config.vocab_size, config.hidden_size, self.padding_idx
-        )
-        self.layers = nn.ModuleList(
-            [
-                LlamaDecoderLayer(config)
-                for _ in range(config.num_hidden_layers)
-            ]
-        )
-        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-        self.gradient_checkpointing = False
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
-    def _prepare_decoder_attention_mask(
-        self,
-        attention_mask,
-        input_shape,
-        inputs_embeds,
-        past_key_values_length,
-    ):
-        # create causal mask
-        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-        combined_attention_mask = None
-        if input_shape[-1] > 1:
-            combined_attention_mask = _make_causal_mask(
-                input_shape,
-                inputs_embeds.dtype,
-                device=inputs_embeds.device,
-                past_key_values_length=past_key_values_length,
-            )
-
-        if attention_mask is not None:
-            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-            expanded_attn_mask = _expand_mask(
-                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
-            ).to(inputs_embeds.device)
-            combined_attention_mask = (
-                expanded_attn_mask
-                if combined_attention_mask is None
-                else expanded_attn_mask + combined_attention_mask
-            )
-
-        return combined_attention_mask
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ):
-        t1 = time()
-        output_attentions = (
-            output_attentions
-            if output_attentions is not None
-            else self.config.output_attentions
-        )
-        output_hidden_states = (
-            output_hidden_states
-            if output_hidden_states is not None
-            else self.config.output_hidden_states
-        )
-        use_cache = (
-            use_cache if use_cache is not None else self.config.use_cache
-        )
-
-        return_dict = (
-            return_dict
-            if return_dict is not None
-            else self.config.use_return_dict
-        )
-
-        # retrieve input_ids and inputs_embeds
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError(
-                "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
-            )
-        elif input_ids is not None:
-            batch_size, seq_length = input_ids.shape
-        elif inputs_embeds is not None:
-            batch_size, seq_length, _ = inputs_embeds.shape
-        else:
-            raise ValueError(
-                "You have to specify either decoder_input_ids or decoder_inputs_embeds"
-            )
-
-        seq_length_with_past = seq_length
-        past_key_values_length = 0
-
-        if past_key_values is not None:
-            past_key_values_length = past_key_values[0][0].shape[2]
-            seq_length_with_past = (
-                seq_length_with_past + past_key_values_length
-            )
-
-        if position_ids is None:
-            device = (
-                input_ids.device
-                if input_ids is not None
-                else inputs_embeds.device
-            )
-            position_ids = torch.arange(
-                past_key_values_length,
-                seq_length + past_key_values_length,
-                dtype=torch.long,
-                device=device,
-            )
-            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
-        else:
-            position_ids = position_ids.view(-1, seq_length).long()
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-        # embed positions
-        if attention_mask is None:
-            attention_mask = torch.ones(
-                (batch_size, seq_length_with_past),
-                dtype=torch.bool,
-                device=inputs_embeds.device,
-            )
-
-        attention_mask = self._prepare_decoder_attention_mask(
-            attention_mask,
-            (batch_size, seq_length),
-            inputs_embeds,
-            past_key_values_length,
-        )
-
-        hidden_states = inputs_embeds
-
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-        next_decoder_cache = () if use_cache else None
-
-        for idx, decoder_layer in enumerate(self.compressedlayers):
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-
-            past_key_value = (
-                past_key_values[8 * idx : 8 * (idx + 1)]
-                if past_key_values is not None
-                else None
-            )
-
-            if self.gradient_checkpointing and self.training:
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        # None for past_key_value
-                        return module(*inputs, output_attentions, None)
-
-                    return custom_forward
-
-                layer_outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(decoder_layer),
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                    None,
-                )
-            else:
-                layer_outputs = decoder_layer.forward(
-                    hidden_states,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    past_key_value=past_key_value,
-                    output_attentions=output_attentions,
-                    use_cache=use_cache,
-                )
-
-            hidden_states = layer_outputs[0]
-
-            if use_cache:
-                next_decoder_cache += (layer_outputs[1:],)
-
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
-
-        try:
-            hidden_states = np.asarray(hidden_states, hidden_states.dtype)
-        except:
-            _ = 10
-
-        hidden_states = self.norm(hidden_states)
-
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        next_cache = next_decoder_cache if use_cache else None
-        next_cache = tuple(itertools.chain.from_iterable(next_cache))
-        print(f"Token generated in {time() - t1} seconds")
-        if not return_dict:
-            return tuple(
-                v
-                for v in [
-                    hidden_states,
-                    next_cache,
-                    all_hidden_states,
-                    all_self_attns,
-                ]
-                if v is not None
-            )
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=next_cache,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-        )
-
-
-class EightLayerLayerSV(torch.nn.Module):
-    def __init__(self, layers):
-        super().__init__()
-        assert len(layers) == 8
-        self.layers = layers
-
-    def forward(
-        self,
-        hidden_states,
-        attention_mask,
-        position_ids,
-        pkv00,
-        pkv01,
-        pkv10,
-        pkv11,
-        pkv20,
-        pkv21,
-        pkv30,
-        pkv31,
-        pkv40,
-        pkv41,
-        pkv50,
-        pkv51,
-        pkv60,
-        pkv61,
-        pkv70,
-        pkv71,
-    ):
-        pkvs = [
-            (pkv00, pkv01),
-            (pkv10, pkv11),
-            (pkv20, pkv21),
-            (pkv30, pkv31),
-            (pkv40, pkv41),
-            (pkv50, pkv51),
-            (pkv60, pkv61),
-            (pkv70, pkv71),
-        ]
-        new_pkvs = []
-        for layer, pkv in zip(self.layers, pkvs):
-            outputs = layer(
-                hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=(
-                    pkv[0],
-                    pkv[1],
-                ),
-                use_cache=True,
-            )
-
-            hidden_states = outputs[0]
-            new_pkvs.append(
-                (
-                    outputs[-1][0],
-                    outputs[-1][1],
-                )
-            )
-        (
-            (new_pkv00, new_pkv01),
-            (new_pkv10, new_pkv11),
-            (new_pkv20, new_pkv21),
-            (new_pkv30, new_pkv31),
-            (new_pkv40, new_pkv41),
-            (new_pkv50, new_pkv51),
-            (new_pkv60, new_pkv61),
-            (new_pkv70, new_pkv71),
-        ) = new_pkvs
-        return (
-            hidden_states,
-            new_pkv00,
-            new_pkv01,
-            new_pkv10,
-            new_pkv11,
-            new_pkv20,
-            new_pkv21,
-            new_pkv30,
-            new_pkv31,
-            new_pkv40,
-            new_pkv41,
-            new_pkv50,
-            new_pkv51,
-            new_pkv60,
-            new_pkv61,
-            new_pkv70,
-            new_pkv71,
-        )
-
-
-class EightLayerLayerFV(torch.nn.Module):
-    def __init__(self, layers):
-        super().__init__()
-        assert len(layers) == 8
-        self.layers = layers
-
-    def forward(self, hidden_states, attention_mask, position_ids):
-        new_pkvs = []
-        for layer in self.layers:
-            outputs = layer(
-                hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=None,
-                use_cache=True,
-            )
-
-            hidden_states = outputs[0]
-            new_pkvs.append(
-                (
-                    outputs[-1][0],
-                    outputs[-1][1],
-                )
-            )
-        (
-            (new_pkv00, new_pkv01),
-            (new_pkv10, new_pkv11),
-            (new_pkv20, new_pkv21),
-            (new_pkv30, new_pkv31),
-            (new_pkv40, new_pkv41),
-            (new_pkv50, new_pkv51),
-            (new_pkv60, new_pkv61),
-            (new_pkv70, new_pkv71),
-        ) = new_pkvs
-        return (
-            hidden_states,
-            new_pkv00,
-            new_pkv01,
-            new_pkv10,
-            new_pkv11,
-            new_pkv20,
-            new_pkv21,
-            new_pkv30,
-            new_pkv31,
-            new_pkv40,
-            new_pkv41,
-            new_pkv50,
-            new_pkv51,
-            new_pkv60,
-            new_pkv61,
-            new_pkv70,
-            new_pkv71,
-        )
-
-
-class CompiledEightLayerLayerSV(torch.nn.Module):
-    def __init__(self, model):
-        super().__init__()
-        self.model = model
-
-    def forward(
-        self,
-        hidden_states,
-        attention_mask,
-        position_ids,
-        past_key_value,
-        output_attentions=False,
-        use_cache=True,
-    ):
-        hidden_states = hidden_states.detach()
-        attention_mask = attention_mask.detach()
-        position_ids = position_ids.detach()
-        (
-            (pkv00, pkv01),
-            (pkv10, pkv11),
-            (pkv20, pkv21),
-            (pkv30, pkv31),
-            (pkv40, pkv41),
-            (pkv50, pkv51),
-            (pkv60, pkv61),
-            (pkv70, pkv71),
-        ) = past_key_value
-        pkv00 = pkv00.detatch()
-        pkv01 = pkv01.detatch()
-        pkv10 = pkv10.detatch()
-        pkv11 = pkv11.detatch()
-        pkv20 = pkv20.detatch()
-        pkv21 = pkv21.detatch()
-        pkv30 = pkv30.detatch()
-        pkv31 = pkv31.detatch()
-        pkv40 = pkv40.detatch()
-        pkv41 = pkv41.detatch()
-        pkv50 = pkv50.detatch()
-        pkv51 = pkv51.detatch()
-        pkv60 = pkv60.detatch()
-        pkv61 = pkv61.detatch()
-        pkv70 = pkv70.detatch()
-        pkv71 = pkv71.detatch()
-
-        output = self.model(
-            "forward",
-            (
-                hidden_states,
-                attention_mask,
-                position_ids,
-                pkv00,
-                pkv01,
-                pkv10,
-                pkv11,
-                pkv20,
-                pkv21,
-                pkv30,
-                pkv31,
-                pkv40,
-                pkv41,
-                pkv50,
-                pkv51,
-                pkv60,
-                pkv61,
-                pkv70,
-                pkv71,
-            ),
-            send_to_host=False,
-        )
-        return (
-            output[0],
-            (output[1][0], output[1][1]),
-            (output[2][0], output[2][1]),
-            (output[3][0], output[3][1]),
-            (output[4][0], output[4][1]),
-            (output[5][0], output[5][1]),
-            (output[6][0], output[6][1]),
-            (output[7][0], output[7][1]),
-            (output[8][0], output[8][1]),
-        )
-
-
-def forward_compressed(
-    self,
-    input_ids: torch.LongTensor = None,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[List[torch.FloatTensor]] = None,
-    inputs_embeds: Optional[torch.FloatTensor] = None,
-    use_cache: Optional[bool] = None,
-    output_attentions: Optional[bool] = None,
-    output_hidden_states: Optional[bool] = None,
-    return_dict: Optional[bool] = None,
-):
-    output_attentions = (
-        output_attentions
-        if output_attentions is not None
-        else self.config.output_attentions
-    )
-    output_hidden_states = (
-        output_hidden_states
-        if output_hidden_states is not None
-        else self.config.output_hidden_states
-    )
-    use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-    return_dict = (
-        return_dict if return_dict is not None else self.config.use_return_dict
-    )
-
-    # retrieve input_ids and inputs_embeds
-    if input_ids is not None and inputs_embeds is not None:
-        raise ValueError(
-            "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
-        )
-    elif input_ids is not None:
-        batch_size, seq_length = input_ids.shape
-    elif inputs_embeds is not None:
-        batch_size, seq_length, _ = inputs_embeds.shape
-    else:
-        raise ValueError(
-            "You have to specify either decoder_input_ids or decoder_inputs_embeds"
-        )
-
-    seq_length_with_past = seq_length
-    past_key_values_length = 0
-
-    if past_key_values is not None:
-        past_key_values_length = past_key_values[0][0].shape[2]
-        seq_length_with_past = seq_length_with_past + past_key_values_length
-
-    if position_ids is None:
-        device = (
-            input_ids.device if input_ids is not None else inputs_embeds.device
-        )
-        position_ids = torch.arange(
-            past_key_values_length,
-            seq_length + past_key_values_length,
-            dtype=torch.long,
-            device=device,
-        )
-        position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
-    else:
-        position_ids = position_ids.view(-1, seq_length).long()
-
-    if inputs_embeds is None:
-        inputs_embeds = self.embed_tokens(input_ids)
-    # embed positions
-    if attention_mask is None:
-        attention_mask = torch.ones(
-            (batch_size, seq_length_with_past),
-            dtype=torch.bool,
-            device=inputs_embeds.device,
-        )
-    attention_mask = self._prepare_decoder_attention_mask(
-        attention_mask,
-        (batch_size, seq_length),
-        inputs_embeds,
-        past_key_values_length,
-    )
-
-    hidden_states = inputs_embeds
-
-    # decoder layers
-    all_hidden_states = () if output_hidden_states else None
-    all_self_attns = () if output_attentions else None
-    next_decoder_cache = () if use_cache else None
-
-    for idx, decoder_layer in enumerate(self.compressedlayers):
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        past_key_value = (
-            past_key_values[8 * idx : 8 * (idx + 1)]
-            if past_key_values is not None
-            else None
-        )
-
-        if self.gradient_checkpointing and self.training:
-
-            def create_custom_forward(module):
-                def custom_forward(*inputs):
-                    # None for past_key_value
-                    return module(*inputs, output_attentions, None)
-
-                return custom_forward
-
-            layer_outputs = torch.utils.checkpoint.checkpoint(
-                create_custom_forward(decoder_layer),
-                hidden_states,
-                attention_mask,
-                position_ids,
-                None,
-            )
-        else:
-            layer_outputs = decoder_layer(
-                hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=past_key_value,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-            )
-
-        hidden_states = layer_outputs[0]
-
-        if use_cache:
-            next_decoder_cache += (
-                layer_outputs[2 if output_attentions else 1],
-            )
-
-        if output_attentions:
-            all_self_attns += (layer_outputs[1],)
-
-    hidden_states = self.norm(hidden_states)
-
-    # add hidden states from the last decoder layer
-    if output_hidden_states:
-        all_hidden_states += (hidden_states,)
-
-    next_cache = next_decoder_cache if use_cache else None
-    if not return_dict:
-        return tuple(
-            v
-            for v in [
-                hidden_states,
-                next_cache,
-                all_hidden_states,
-                all_self_attns,
-            ]
-            if v is not None
-        )
-    return BaseModelOutputWithPast(
-        last_hidden_state=hidden_states,
-        past_key_values=next_cache,
-        hidden_states=all_hidden_states,
-        attentions=all_self_attns,
-    )
-
-
-class CompiledEightLayerLayer(torch.nn.Module):
-    def __init__(self, model):
-        super().__init__()
-        self.model = model
-
-    def forward(
-        self,
-        hidden_states,
-        attention_mask,
-        position_ids,
-        past_key_value=None,
-        output_attentions=False,
-        use_cache=True,
-    ):
-        t2 = time()
-        if past_key_value is None:
-            try:
-                hidden_states = np.asarray(hidden_states, hidden_states.dtype)
-            except:
-                pass
-            attention_mask = attention_mask.detach()
-            position_ids = position_ids.detach()
-            t1 = time()
-
-            output = self.model(
-                "first_vicuna_forward",
-                (hidden_states, attention_mask, position_ids),
-                send_to_host=False,
-            )
-            output2 = (
-                output[0],
-                (
-                    output[1],
-                    output[2],
-                ),
-                (
-                    output[3],
-                    output[4],
-                ),
-                (
-                    output[5],
-                    output[6],
-                ),
-                (
-                    output[7],
-                    output[8],
-                ),
-                (
-                    output[9],
-                    output[10],
-                ),
-                (
-                    output[11],
-                    output[12],
-                ),
-                (
-                    output[13],
-                    output[14],
-                ),
-                (
-                    output[15],
-                    output[16],
-                ),
-            )
-            return output2
-        else:
-            (
-                (pkv00, pkv01),
-                (pkv10, pkv11),
-                (pkv20, pkv21),
-                (pkv30, pkv31),
-                (pkv40, pkv41),
-                (pkv50, pkv51),
-                (pkv60, pkv61),
-                (pkv70, pkv71),
-            ) = past_key_value
-
-            try:
-                hidden_states = hidden_states.detach()
-                attention_mask = attention_mask.detach()
-                position_ids = position_ids.detach()
-                pkv00 = pkv00.detach()
-                pkv01 = pkv01.detach()
-                pkv10 = pkv10.detach()
-                pkv11 = pkv11.detach()
-                pkv20 = pkv20.detach()
-                pkv21 = pkv21.detach()
-                pkv30 = pkv30.detach()
-                pkv31 = pkv31.detach()
-                pkv40 = pkv40.detach()
-                pkv41 = pkv41.detach()
-                pkv50 = pkv50.detach()
-                pkv51 = pkv51.detach()
-                pkv60 = pkv60.detach()
-                pkv61 = pkv61.detach()
-                pkv70 = pkv70.detach()
-                pkv71 = pkv71.detach()
-            except:
-                x = 10
-
-            t1 = time()
-            if type(hidden_states) == iree.runtime.array_interop.DeviceArray:
-                hidden_states = np.array(hidden_states, hidden_states.dtype)
-                hidden_states = torch.tensor(hidden_states)
-                hidden_states = hidden_states.detach()
-
-            output = self.model(
-                "second_vicuna_forward",
-                (
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                    pkv00,
-                    pkv01,
-                    pkv10,
-                    pkv11,
-                    pkv20,
-                    pkv21,
-                    pkv30,
-                    pkv31,
-                    pkv40,
-                    pkv41,
-                    pkv50,
-                    pkv51,
-                    pkv60,
-                    pkv61,
-                    pkv70,
-                    pkv71,
-                ),
-                send_to_host=False,
-            )
-            print(f"{time() - t1}")
-            del pkv00
-            del pkv01
-            del pkv10
-            del pkv11
-            del pkv20
-            del pkv21
-            del pkv30
-            del pkv31
-            del pkv40
-            del pkv41
-            del pkv50
-            del pkv51
-            del pkv60
-            del pkv61
-            del pkv70
-            del pkv71
-            output2 = (
-                output[0],
-                (
-                    output[1],
-                    output[2],
-                ),
-                (
-                    output[3],
-                    output[4],
-                ),
-                (
-                    output[5],
-                    output[6],
-                ),
-                (
-                    output[7],
-                    output[8],
-                ),
-                (
-                    output[9],
-                    output[10],
-                ),
-                (
-                    output[11],
-                    output[12],
-                ),
-                (
-                    output[13],
-                    output[14],
-                ),
-                (
-                    output[15],
-                    output[16],
-                ),
-            )
-            return output2

apps/language_models/src/model_wrappers/vicuna_model.py

@@ -3,606 +3,16 @@ from transformers import AutoModelForCausalLM
 
 
 class FirstVicuna(torch.nn.Module):
-    def __init__(
-        self,
-        model_path,
-        precision="fp32",
-        accumulates="fp32",
-        weight_group_size=128,
-        model_name="vicuna",
-        hf_auth_token: str = None,
-    ):
+    def __init__(self, model_path):
         super().__init__()
         kwargs = {"torch_dtype": torch.float32}
-        if "llama2" in model_name:
-            kwargs["use_auth_token"] = hf_auth_token
-        self.accumulates = (
-            torch.float32 if accumulates == "fp32" else torch.float16
-        )
         self.model = AutoModelForCausalLM.from_pretrained(
             model_path, low_cpu_mem_usage=True, **kwargs
         )
-        print(f"[DEBUG] model_path : {model_path}")
-        if precision in ["int4", "int8"]:
-            from brevitas_examples.common.generative.quantize import (
-                quantize_model,
-            )
-            from brevitas_examples.llm.llm_quant.run_utils import (
-                get_model_impl,
-            )
-
-            print("First Vicuna applying weight quantization..")
-            weight_bit_width = 4 if precision == "int4" else 8
-            quantize_model(
-                get_model_impl(self.model).layers,
-                dtype=self.accumulates,
-                weight_bit_width=weight_bit_width,
-                weight_param_method="stats",
-                weight_scale_precision="float_scale",
-                weight_quant_type="asym",
-                weight_quant_granularity="per_group",
-                weight_group_size=weight_group_size,
-                quantize_weight_zero_point=False,
-            )
-            print("Weight quantization applied.")
 
     def forward(self, input_ids):
         op = self.model(input_ids=input_ids, use_cache=True)
         return_vals = []
-        token = torch.argmax(op.logits[:, -1, :], dim=1)
-        return_vals.append(token)
-
-        temp_past_key_values = op.past_key_values
-        for item in temp_past_key_values:
-            return_vals.append(item[0])
-            return_vals.append(item[1])
-        return tuple(return_vals)
-
-
-class SecondVicuna7B(torch.nn.Module):
-    def __init__(
-        self,
-        model_path,
-        precision="fp32",
-        accumulates="fp32",
-        weight_group_size=128,
-        model_name="vicuna",
-        hf_auth_token: str = None,
-    ):
-        super().__init__()
-        kwargs = {"torch_dtype": torch.float32}
-        if "llama2" in model_name:
-            kwargs["use_auth_token"] = hf_auth_token
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_path, low_cpu_mem_usage=True, **kwargs
-        )
-        self.accumulates = (
-            torch.float32 if accumulates == "fp32" else torch.float16
-        )
-        print(f"[DEBUG] model_path : {model_path}")
-        if precision in ["int4", "int8"]:
-            from brevitas_examples.common.generative.quantize import (
-                quantize_model,
-            )
-            from brevitas_examples.llm.llm_quant.run_utils import (
-                get_model_impl,
-            )
-
-            print("Second Vicuna applying weight quantization..")
-            weight_bit_width = 4 if precision == "int4" else 8
-            quantize_model(
-                get_model_impl(self.model).layers,
-                dtype=self.accumulates,
-                weight_bit_width=weight_bit_width,
-                weight_param_method="stats",
-                weight_scale_precision="float_scale",
-                weight_quant_type="asym",
-                weight_quant_granularity="per_group",
-                weight_group_size=weight_group_size,
-                quantize_weight_zero_point=False,
-            )
-            print("Weight quantization applied.")
-
-    def forward(
-        self,
-        i0,
-        i1,
-        i2,
-        i3,
-        i4,
-        i5,
-        i6,
-        i7,
-        i8,
-        i9,
-        i10,
-        i11,
-        i12,
-        i13,
-        i14,
-        i15,
-        i16,
-        i17,
-        i18,
-        i19,
-        i20,
-        i21,
-        i22,
-        i23,
-        i24,
-        i25,
-        i26,
-        i27,
-        i28,
-        i29,
-        i30,
-        i31,
-        i32,
-        i33,
-        i34,
-        i35,
-        i36,
-        i37,
-        i38,
-        i39,
-        i40,
-        i41,
-        i42,
-        i43,
-        i44,
-        i45,
-        i46,
-        i47,
-        i48,
-        i49,
-        i50,
-        i51,
-        i52,
-        i53,
-        i54,
-        i55,
-        i56,
-        i57,
-        i58,
-        i59,
-        i60,
-        i61,
-        i62,
-        i63,
-        i64,
-    ):
-        token = i0
-        past_key_values = (
-            (i1, i2),
-            (
-                i3,
-                i4,
-            ),
-            (
-                i5,
-                i6,
-            ),
-            (
-                i7,
-                i8,
-            ),
-            (
-                i9,
-                i10,
-            ),
-            (
-                i11,
-                i12,
-            ),
-            (
-                i13,
-                i14,
-            ),
-            (
-                i15,
-                i16,
-            ),
-            (
-                i17,
-                i18,
-            ),
-            (
-                i19,
-                i20,
-            ),
-            (
-                i21,
-                i22,
-            ),
-            (
-                i23,
-                i24,
-            ),
-            (
-                i25,
-                i26,
-            ),
-            (
-                i27,
-                i28,
-            ),
-            (
-                i29,
-                i30,
-            ),
-            (
-                i31,
-                i32,
-            ),
-            (
-                i33,
-                i34,
-            ),
-            (
-                i35,
-                i36,
-            ),
-            (
-                i37,
-                i38,
-            ),
-            (
-                i39,
-                i40,
-            ),
-            (
-                i41,
-                i42,
-            ),
-            (
-                i43,
-                i44,
-            ),
-            (
-                i45,
-                i46,
-            ),
-            (
-                i47,
-                i48,
-            ),
-            (
-                i49,
-                i50,
-            ),
-            (
-                i51,
-                i52,
-            ),
-            (
-                i53,
-                i54,
-            ),
-            (
-                i55,
-                i56,
-            ),
-            (
-                i57,
-                i58,
-            ),
-            (
-                i59,
-                i60,
-            ),
-            (
-                i61,
-                i62,
-            ),
-            (
-                i63,
-                i64,
-            ),
-        )
-        op = self.model(
-            input_ids=token, use_cache=True, past_key_values=past_key_values
-        )
-        return_vals = []
-        token = torch.argmax(op.logits[:, -1, :], dim=1)
-        return_vals.append(token)
-        temp_past_key_values = op.past_key_values
-        for item in temp_past_key_values:
-            return_vals.append(item[0])
-            return_vals.append(item[1])
-        return tuple(return_vals)
-
-
-class SecondVicuna13B(torch.nn.Module):
-    def __init__(
-        self,
-        model_path,
-        precision="int8",
-        accumulates="fp32",
-        weight_group_size=128,
-        model_name="vicuna",
-        hf_auth_token: str = None,
-    ):
-        super().__init__()
-        kwargs = {"torch_dtype": torch.float32}
-        if "llama2" in model_name:
-            kwargs["use_auth_token"] = hf_auth_token
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_path, low_cpu_mem_usage=True, **kwargs
-        )
-        self.accumulates = (
-            torch.float32 if accumulates == "fp32" else torch.float16
-        )
-        if precision in ["int4", "int8"]:
-            from brevitas_examples.common.generative.quantize import (
-                quantize_model,
-            )
-            from brevitas_examples.llm.llm_quant.run_utils import (
-                get_model_impl,
-            )
-
-            print("Second Vicuna applying weight quantization..")
-            weight_bit_width = 4 if precision == "int4" else 8
-            quantize_model(
-                get_model_impl(self.model).layers,
-                dtype=self.accumulates,
-                weight_bit_width=weight_bit_width,
-                weight_param_method="stats",
-                weight_scale_precision="float_scale",
-                weight_quant_type="asym",
-                weight_quant_granularity="per_group",
-                weight_group_size=weight_group_size,
-                quantize_weight_zero_point=False,
-            )
-            print("Weight quantization applied.")
-
-    def forward(
-        self,
-        i0,
-        i1,
-        i2,
-        i3,
-        i4,
-        i5,
-        i6,
-        i7,
-        i8,
-        i9,
-        i10,
-        i11,
-        i12,
-        i13,
-        i14,
-        i15,
-        i16,
-        i17,
-        i18,
-        i19,
-        i20,
-        i21,
-        i22,
-        i23,
-        i24,
-        i25,
-        i26,
-        i27,
-        i28,
-        i29,
-        i30,
-        i31,
-        i32,
-        i33,
-        i34,
-        i35,
-        i36,
-        i37,
-        i38,
-        i39,
-        i40,
-        i41,
-        i42,
-        i43,
-        i44,
-        i45,
-        i46,
-        i47,
-        i48,
-        i49,
-        i50,
-        i51,
-        i52,
-        i53,
-        i54,
-        i55,
-        i56,
-        i57,
-        i58,
-        i59,
-        i60,
-        i61,
-        i62,
-        i63,
-        i64,
-        i65,
-        i66,
-        i67,
-        i68,
-        i69,
-        i70,
-        i71,
-        i72,
-        i73,
-        i74,
-        i75,
-        i76,
-        i77,
-        i78,
-        i79,
-        i80,
-    ):
-        token = i0
-        past_key_values = (
-            (i1, i2),
-            (
-                i3,
-                i4,
-            ),
-            (
-                i5,
-                i6,
-            ),
-            (
-                i7,
-                i8,
-            ),
-            (
-                i9,
-                i10,
-            ),
-            (
-                i11,
-                i12,
-            ),
-            (
-                i13,
-                i14,
-            ),
-            (
-                i15,
-                i16,
-            ),
-            (
-                i17,
-                i18,
-            ),
-            (
-                i19,
-                i20,
-            ),
-            (
-                i21,
-                i22,
-            ),
-            (
-                i23,
-                i24,
-            ),
-            (
-                i25,
-                i26,
-            ),
-            (
-                i27,
-                i28,
-            ),
-            (
-                i29,
-                i30,
-            ),
-            (
-                i31,
-                i32,
-            ),
-            (
-                i33,
-                i34,
-            ),
-            (
-                i35,
-                i36,
-            ),
-            (
-                i37,
-                i38,
-            ),
-            (
-                i39,
-                i40,
-            ),
-            (
-                i41,
-                i42,
-            ),
-            (
-                i43,
-                i44,
-            ),
-            (
-                i45,
-                i46,
-            ),
-            (
-                i47,
-                i48,
-            ),
-            (
-                i49,
-                i50,
-            ),
-            (
-                i51,
-                i52,
-            ),
-            (
-                i53,
-                i54,
-            ),
-            (
-                i55,
-                i56,
-            ),
-            (
-                i57,
-                i58,
-            ),
-            (
-                i59,
-                i60,
-            ),
-            (
-                i61,
-                i62,
-            ),
-            (
-                i63,
-                i64,
-            ),
-            (
-                i65,
-                i66,
-            ),
-            (
-                i67,
-                i68,
-            ),
-            (
-                i69,
-                i70,
-            ),
-            (
-                i71,
-                i72,
-            ),
-            (
-                i73,
-                i74,
-            ),
-            (
-                i75,
-                i76,
-            ),
-            (
-                i77,
-                i78,
-            ),
-            (
-                i79,
-                i80,
-            ),
-        )
-        op = self.model(
-            input_ids=token, use_cache=True, past_key_values=past_key_values
-        )
-        return_vals = []
         return_vals.append(op.logits)
         temp_past_key_values = op.past_key_values
         for item in temp_past_key_values:
@@ -611,49 +21,13 @@ class SecondVicuna13B(torch.nn.Module):
         return tuple(return_vals)
 
 
-class SecondVicuna70B(torch.nn.Module):
-    def __init__(
-        self,
-        model_path,
-        precision="fp32",
-        accumulates="fp32",
-        weight_group_size=128,
-        model_name="vicuna",
-        hf_auth_token: str = None,
-    ):
+class SecondVicuna(torch.nn.Module):
+    def __init__(self, model_path):
         super().__init__()
         kwargs = {"torch_dtype": torch.float32}
-        if "llama2" in model_name:
-            kwargs["use_auth_token"] = hf_auth_token
         self.model = AutoModelForCausalLM.from_pretrained(
             model_path, low_cpu_mem_usage=True, **kwargs
         )
-        self.accumulates = (
-            torch.float32 if accumulates == "fp32" else torch.float16
-        )
-        print(f"[DEBUG] model_path : {model_path}")
-        if precision in ["int4", "int8"]:
-            from brevitas_examples.common.generative.quantize import (
-                quantize_model,
-            )
-            from brevitas_examples.llm.llm_quant.run_utils import (
-                get_model_impl,
-            )
-
-            print("Second Vicuna applying weight quantization..")
-            weight_bit_width = 4 if precision == "int4" else 8
-            quantize_model(
-                get_model_impl(self.model).layers,
-                dtype=self.accumulates,
-                weight_bit_width=weight_bit_width,
-                weight_param_method="stats",
-                weight_scale_precision="float_scale",
-                weight_quant_type="asym",
-                weight_quant_granularity="per_group",
-                weight_group_size=weight_group_size,
-                quantize_weight_zero_point=False,
-            )
-            print("Weight quantization applied.")
 
     def forward(
         self,
@@ -722,103 +96,9 @@ class SecondVicuna70B(torch.nn.Module):
         i62,
         i63,
         i64,
-        i65,
-        i66,
-        i67,
-        i68,
-        i69,
-        i70,
-        i71,
-        i72,
-        i73,
-        i74,
-        i75,
-        i76,
-        i77,
-        i78,
-        i79,
-        i80,
-        i81,
-        i82,
-        i83,
-        i84,
-        i85,
-        i86,
-        i87,
-        i88,
-        i89,
-        i90,
-        i91,
-        i92,
-        i93,
-        i94,
-        i95,
-        i96,
-        i97,
-        i98,
-        i99,
-        i100,
-        i101,
-        i102,
-        i103,
-        i104,
-        i105,
-        i106,
-        i107,
-        i108,
-        i109,
-        i110,
-        i111,
-        i112,
-        i113,
-        i114,
-        i115,
-        i116,
-        i117,
-        i118,
-        i119,
-        i120,
-        i121,
-        i122,
-        i123,
-        i124,
-        i125,
-        i126,
-        i127,
-        i128,
-        i129,
-        i130,
-        i131,
-        i132,
-        i133,
-        i134,
-        i135,
-        i136,
-        i137,
-        i138,
-        i139,
-        i140,
-        i141,
-        i142,
-        i143,
-        i144,
-        i145,
-        i146,
-        i147,
-        i148,
-        i149,
-        i150,
-        i151,
-        i152,
-        i153,
-        i154,
-        i155,
-        i156,
-        i157,
-        i158,
-        i159,
-        i160,
     ):
+        # input_ids = input_tuple[0]
+        # input_tuple = torch.unbind(pkv, dim=0)
         token = i0
         past_key_values = (
             (i1, i2),
@@ -946,198 +226,6 @@ class SecondVicuna70B(torch.nn.Module):
                 i63,
                 i64,
             ),
-            (
-                i65,
-                i66,
-            ),
-            (
-                i67,
-                i68,
-            ),
-            (
-                i69,
-                i70,
-            ),
-            (
-                i71,
-                i72,
-            ),
-            (
-                i73,
-                i74,
-            ),
-            (
-                i75,
-                i76,
-            ),
-            (
-                i77,
-                i78,
-            ),
-            (
-                i79,
-                i80,
-            ),
-            (
-                i81,
-                i82,
-            ),
-            (
-                i83,
-                i84,
-            ),
-            (
-                i85,
-                i86,
-            ),
-            (
-                i87,
-                i88,
-            ),
-            (
-                i89,
-                i90,
-            ),
-            (
-                i91,
-                i92,
-            ),
-            (
-                i93,
-                i94,
-            ),
-            (
-                i95,
-                i96,
-            ),
-            (
-                i97,
-                i98,
-            ),
-            (
-                i99,
-                i100,
-            ),
-            (
-                i101,
-                i102,
-            ),
-            (
-                i103,
-                i104,
-            ),
-            (
-                i105,
-                i106,
-            ),
-            (
-                i107,
-                i108,
-            ),
-            (
-                i109,
-                i110,
-            ),
-            (
-                i111,
-                i112,
-            ),
-            (
-                i113,
-                i114,
-            ),
-            (
-                i115,
-                i116,
-            ),
-            (
-                i117,
-                i118,
-            ),
-            (
-                i119,
-                i120,
-            ),
-            (
-                i121,
-                i122,
-            ),
-            (
-                i123,
-                i124,
-            ),
-            (
-                i125,
-                i126,
-            ),
-            (
-                i127,
-                i128,
-            ),
-            (
-                i129,
-                i130,
-            ),
-            (
-                i131,
-                i132,
-            ),
-            (
-                i133,
-                i134,
-            ),
-            (
-                i135,
-                i136,
-            ),
-            (
-                i137,
-                i138,
-            ),
-            (
-                i139,
-                i140,
-            ),
-            (
-                i141,
-                i142,
-            ),
-            (
-                i143,
-                i144,
-            ),
-            (
-                i145,
-                i146,
-            ),
-            (
-                i147,
-                i148,
-            ),
-            (
-                i149,
-                i150,
-            ),
-            (
-                i151,
-                i152,
-            ),
-            (
-                i153,
-                i154,
-            ),
-            (
-                i155,
-                i156,
-            ),
-            (
-                i157,
-                i158,
-            ),
-            (
-                i159,
-                i160,
-            ),
         )
         op = self.model(
             input_ids=token, use_cache=True, past_key_values=past_key_values
@@ -1159,17 +247,15 @@ class CombinedModel(torch.nn.Module):
     ):
         super().__init__()
         self.first_vicuna = FirstVicuna(first_vicuna_model_path)
-        # NOT using this path for 13B currently, hence using `SecondVicuna7B`.
-        self.second_vicuna = SecondVicuna7B(second_vicuna_model_path)
+        self.second_vicuna = SecondVicuna(second_vicuna_model_path)
 
     def forward(self, input_ids):
-        first_output = self.first_vicuna(input_ids=input_ids)
-        # generate second vicuna
-        compilation_input_ids = torch.zeros([1, 1], dtype=torch.int64)
-        pkv = tuple(
-            (torch.zeros([1, 32, 19, 128], dtype=torch.float32))
-            for _ in range(64)
-        )
-        secondVicunaCompileInput = (compilation_input_ids,) + pkv
-        second_output = self.second_vicuna(*secondVicunaCompileInput)
+        first_output = self.first_vicuna(input_ids=input_ids, use_cache=True)
+        logits = first_output[0]
+        pkv = first_output[1:]
+
+        token = torch.argmax(torch.tensor(logits)[:, -1, :], dim=1)
+        token = token.to(torch.int64).reshape([1, 1])
+        secondVicunaInput = (token,) + tuple(pkv)
+        second_output = self.second_vicuna(secondVicunaInput)
         return second_output

apps/language_models/src/model_wrappers/vicuna_sharded_model.py

@@ -1,5 +1,4 @@
 import torch
-import time
 
 
 class FirstVicunaLayer(torch.nn.Module):
@@ -63,117 +62,11 @@ class SecondVicunaLayer(torch.nn.Module):
         )
 
 
-class ShardedVicunaModel(torch.nn.Module):
-    def __init__(self, model, layers, lmhead, embedding, norm):
-        super().__init__()
-        self.model = model
-        self.model.model.config.use_cache = True
-        self.model.model.config.output_attentions = False
-        self.layers = layers
-        self.norm = norm
-        self.embedding = embedding
-        self.lmhead = lmhead
-        self.model.model.norm = self.norm
-        self.model.model.embed_tokens = self.embedding
-        self.model.lm_head = self.lmhead
-        self.model.model.layers = torch.nn.modules.container.ModuleList(
-            self.layers
-        )
-
-    def forward(
-        self,
-        input_ids,
-        is_first=True,
-        past_key_values=None,
-        attention_mask=None,
-    ):
-        return self.model.forward(
-            input_ids,
-            attention_mask=attention_mask,
-            past_key_values=past_key_values,
-        )
-
-
-class LMHead(torch.nn.Module):
-    def __init__(self, model):
-        super().__init__()
-        self.model = model
-
-    def forward(self, hidden_states):
-        output = self.model(hidden_states)
-        return output
-
-
-class LMHeadCompiled(torch.nn.Module):
+class CompiledFirstVicunaLayer(torch.nn.Module):
     def __init__(self, shark_module):
         super().__init__()
         self.model = shark_module
 
-    def forward(self, hidden_states):
-        hidden_states_sample = hidden_states.detach()
-
-        output = self.model("forward", (hidden_states,))
-        output = torch.tensor(output)
-
-        return output
-
-
-class VicunaNorm(torch.nn.Module):
-    def __init__(self, model):
-        super().__init__()
-        self.model = model
-
-    def forward(self, hidden_states):
-        output = self.model(hidden_states)
-        return output
-
-
-class VicunaNormCompiled(torch.nn.Module):
-    def __init__(self, shark_module):
-        super().__init__()
-        self.model = shark_module
-
-    def forward(self, hidden_states):
-        try:
-            hidden_states.detach()
-        except:
-            pass
-        output = self.model("forward", (hidden_states,), send_to_host=True)
-        output = torch.tensor(output)
-
-        return output
-
-
-class VicunaEmbedding(torch.nn.Module):
-    def __init__(self, model):
-        super().__init__()
-        self.model = model
-
-    def forward(self, input_ids):
-        output = self.model(input_ids)
-        return output
-
-
-class VicunaEmbeddingCompiled(torch.nn.Module):
-    def __init__(self, shark_module):
-        super().__init__()
-        self.model = shark_module
-
-    def forward(self, input_ids):
-        input_ids.detach()
-        output = self.model("forward", (input_ids,), send_to_host=True)
-        output = torch.tensor(output)
-
-        return output
-
-
-class CompiledVicunaLayer(torch.nn.Module):
-    def __init__(self, shark_module, idx, breakpoints):
-        super().__init__()
-        self.model = shark_module
-        self.idx = idx
-        self.breakpoints = breakpoints
-
     def forward(
         self,
         hidden_states,
@@ -183,65 +76,103 @@ class CompiledVicunaLayer(torch.nn.Module):
         output_attentions=False,
         use_cache=True,
     ):
-        if self.breakpoints is None:
-            is_breakpoint = False
+        hidden_states = hidden_states.detach()
+        attention_mask = attention_mask.detach()
+        position_ids = position_ids.detach()
+        output = self.model(
+            "forward",
+            (
+                hidden_states,
+                attention_mask,
+                position_ids,
+            ),
+        )
+
+        output0 = torch.tensor(output[0])
+        output1 = torch.tensor(output[1])
+        output2 = torch.tensor(output[2])
+
+        return (
+            output0,
+            (
+                output1,
+                output2,
+            ),
+        )
+
+
+class CompiledSecondVicunaLayer(torch.nn.Module):
+    def __init__(self, shark_module):
+        super().__init__()
+        self.model = shark_module
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        position_ids,
+        past_key_value,
+        output_attentions=False,
+        use_cache=True,
+    ):
+        hidden_states = hidden_states.detach()
+        attention_mask = attention_mask.detach()
+        position_ids = position_ids.detach()
+        pkv0 = past_key_value[0].detach()
+        pkv1 = past_key_value[1].detach()
+        output = self.model(
+            "forward",
+            (
+                hidden_states,
+                attention_mask,
+                position_ids,
+                pkv0,
+                pkv1,
+            ),
+        )
+
+        output0 = torch.tensor(output[0])
+        output1 = torch.tensor(output[1])
+        output2 = torch.tensor(output[2])
+
+        return (
+            output0,
+            (
+                output1,
+                output2,
+            ),
+        )
+
+
+class ShardedVicunaModel(torch.nn.Module):
+    def __init__(self, model, layers0, layers1):
+        super().__init__()
+        self.model = model
+        assert len(layers0) == len(model.model.layers)
+        # self.model.model.layers = torch.nn.modules.container.ModuleList(layers0)
+        self.model.model.config.use_cache = True
+        self.model.model.config.output_attentions = False
+        self.layers0 = layers0
+        self.layers1 = layers1
+
+    def forward(
+        self,
+        input_ids,
+        is_first=True,
+        past_key_values=None,
+        attention_mask=None,
+    ):
+        if is_first:
+            self.model.model.layers = torch.nn.modules.container.ModuleList(
+                self.layers0
+            )
+            return self.model.forward(input_ids, attention_mask=attention_mask)
         else:
-            is_breakpoint = self.idx + 1 in self.breakpoints
-        if past_key_value is None:
-            output = self.model(
-                "first_vicuna_forward",
-                (
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                ),
-                send_to_host=is_breakpoint,
+            self.model.model.layers = torch.nn.modules.container.ModuleList(
+                self.layers1
             )
-
-            if is_breakpoint:
-                output0 = torch.tensor(output[0])
-                output1 = torch.tensor(output[1])
-                output2 = torch.tensor(output[2])
-            else:
-                output0 = output[0]
-                output1 = output[1]
-                output2 = output[2]
-
-            return (
-                output0,
-                (
-                    output1,
-                    output2,
-                ),
-            )
-        else:
-            pkv0 = past_key_value[0]
-            pkv1 = past_key_value[1]
-            output = self.model(
-                "second_vicuna_forward",
-                (
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                    pkv0,
-                    pkv1,
-                ),
-                send_to_host=is_breakpoint,
-            )
-
-            if is_breakpoint:
-                output0 = torch.tensor(output[0])
-                output1 = torch.tensor(output[1])
-                output2 = torch.tensor(output[2])
-            else:
-                output0 = output[0]
-                output1 = output[1]
-                output2 = output[2]
-
-            return (
-                output0,
-                (
-                    output1,
-                    output2,
-                ),
+            return self.model.forward(
+                input_ids,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
             )

apps/language_models/src/pipelines/SharkLLMBase.py

@@ -3,10 +3,7 @@ from abc import ABC, abstractmethod
 
 class SharkLLMBase(ABC):
     def __init__(
-        self,
-        model_name,
-        hf_model_path=None,
-        max_num_tokens=512,
+        self, model_name, hf_model_path=None, max_num_tokens=512
     ) -> None:
         self.model_name = model_name
         self.hf_model_path = hf_model_path

apps/language_models/src/pipelines/falcon_pipeline.py

@@ -1,17 +1,4 @@
 from apps.language_models.src.model_wrappers.falcon_model import FalconModel
-from apps.language_models.src.model_wrappers.falcon_sharded_model import (
-    WordEmbeddingsLayer,
-    CompiledWordEmbeddingsLayer,
-    LNFEmbeddingLayer,
-    CompiledLNFEmbeddingLayer,
-    LMHeadEmbeddingLayer,
-    CompiledLMHeadEmbeddingLayer,
-    FourWayShardingDecoderLayer,
-    TwoWayShardingDecoderLayer,
-    CompiledFourWayShardingDecoderLayer,
-    CompiledTwoWayShardingDecoderLayer,
-    ShardedFalconModel,
-)
 from apps.language_models.src.pipelines.SharkLLMBase import SharkLLMBase
 from apps.language_models.utils import (
     get_vmfb_from_path,
@@ -20,39 +7,30 @@ from io import BytesIO
 from pathlib import Path
 from contextlib import redirect_stdout
 from shark.shark_downloader import download_public_file
-from shark.shark_importer import import_with_fx, save_mlir
+from shark.shark_importer import import_with_fx
 from shark.shark_inference import SharkInference
-from transformers import AutoTokenizer, AutoModelForCausalLM, GPTQConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from transformers.generation import (
     GenerationConfig,
     LogitsProcessorList,
     StoppingCriteriaList,
 )
 import copy
-import time
+
 import re
 import torch
 import torch_mlir
 import os
 import argparse
-import gc
 
 parser = argparse.ArgumentParser(
     prog="falcon runner",
     description="runs a falcon model",
 )
 
+parser.add_argument("--falcon_variant_to_use", default="7b", help="7b, 40b")
 parser.add_argument(
-    "--falcon_variant_to_use", default="7b", help="7b, 40b, 180b"
-)
-parser.add_argument(
-    "--compressed",
-    default=False,
-    action=argparse.BooleanOptionalAction,
-    help="Do the compression of sharded layers",
-)
-parser.add_argument(
-    "--precision", "-p", default="fp16", choices=["fp32", "fp16", "int4"]
+    "--precision", "-p", default="fp16", help="fp32, fp16, int8, int4"
 )
 parser.add_argument("--device", "-d", default="cuda", help="vulkan, cpu, cuda")
 parser.add_argument(
@@ -71,7 +49,7 @@ parser.add_argument(
 )
 parser.add_argument(
     "--load_mlir_from_shark_tank",
-    default=True,
+    default=False,
     action=argparse.BooleanOptionalAction,
     help="download precompile mlir from shark tank",
 )
@@ -81,74 +59,32 @@ parser.add_argument(
     action=argparse.BooleanOptionalAction,
     help="Run model in cli mode",
 )
-parser.add_argument(
-    "--hf_auth_token",
-    type=str,
-    default=None,
-    help="Specify your own huggingface authentication token for falcon-180B model.",
-)
-parser.add_argument(
-    "-s",
-    "--sharded",
-    default=False,
-    action=argparse.BooleanOptionalAction,
-    help="Run model as sharded",
-)
-parser.add_argument(
-    "--num_shards",
-    type=int,
-    default=4,
-    choices=[2, 4],
-    help="Number of shards.",
-)
 
 
-class ShardedFalcon(SharkLLMBase):
+class Falcon(SharkLLMBase):
     def __init__(
         self,
         model_name,
-        hf_model_path="tiiuae/falcon-7b-instruct",
-        hf_auth_token: str = None,
+        hf_model_path,
         max_num_tokens=150,
         device="cuda",
         precision="fp32",
         falcon_mlir_path=None,
         falcon_vmfb_path=None,
-        debug=False,
     ) -> None:
         super().__init__(model_name, hf_model_path, max_num_tokens)
-        print("hf_model_path: ", self.hf_model_path)
-
-        if (
-            "180b" in self.model_name
-            and precision != "int4"
-            and hf_auth_token == None
-        ):
-            raise ValueError(
-                """ HF auth token required for falcon-180b. Pass it using
-                --hf_auth_token flag. You can ask for the access to the model
-                here: https://huggingface.co/tiiuae/falcon-180B-chat."""
-            )
-
-        if args.sharded and "180b" not in self.model_name:
-            raise ValueError("Sharding supported only for Falcon-180B")
-
-        self.hf_auth_token = hf_auth_token
         self.max_padding_length = 100
         self.device = device
         self.precision = precision
         self.falcon_vmfb_path = falcon_vmfb_path
         self.falcon_mlir_path = falcon_mlir_path
-        self.debug = debug
         self.tokenizer = self.get_tokenizer()
-        self.src_model = self.get_src_model()
         self.shark_model = self.compile()
+        self.src_model = self.get_src_model()
 
     def get_tokenizer(self):
         tokenizer = AutoTokenizer.from_pretrained(
-            self.hf_model_path,
-            trust_remote_code=True,
-            token=self.hf_auth_token,
+            self.hf_model_path, trust_remote_code=True
         )
         tokenizer.padding_side = "left"
         tokenizer.pad_token_id = 11
@@ -156,535 +92,13 @@ class ShardedFalcon(SharkLLMBase):
 
     def get_src_model(self):
         print("Loading src model: ", self.model_name)
-        kwargs = {
-            "torch_dtype": torch.float32,
-            "trust_remote_code": True,
-            "token": self.hf_auth_token,
-        }
-        if self.precision == "int4":
-            quantization_config = GPTQConfig(bits=4, disable_exllama=True)
-            kwargs["quantization_config"] = quantization_config
-            kwargs["device_map"] = "cpu"
+        kwargs = {"torch_dtype": torch.float, "trust_remote_code": True}
         falcon_model = AutoModelForCausalLM.from_pretrained(
             self.hf_model_path, **kwargs
         )
         return falcon_model
 
-    def compile_layer(
-        self, layer, falconCompileInput, layer_id, device_idx=None
-    ):
-        self.falcon_mlir_path = Path(
-            f"falcon_{args.falcon_variant_to_use}_layer_{layer_id}_{self.precision}.mlir"
-        )
-        self.falcon_vmfb_path = Path(
-            f"falcon_{args.falcon_variant_to_use}_layer_{layer_id}_{self.precision}_{self.device}.vmfb"
-        )
-
-        if args.use_precompiled_model:
-            if not self.falcon_vmfb_path.exists():
-                # Downloading VMFB from shark_tank
-                print(f"[DEBUG] Trying to download vmfb from shark_tank")
-                download_public_file(
-                    f"gs://shark_tank/falcon/sharded/falcon_{args.falcon_variant_to_use}/vmfb/"
-                    + str(self.falcon_vmfb_path),
-                    self.falcon_vmfb_path.absolute(),
-                    single_file=True,
-                )
-            vmfb = get_vmfb_from_path(
-                self.falcon_vmfb_path,
-                self.device,
-                "linalg",
-                device_id=device_idx,
-            )
-            if vmfb is not None:
-                return vmfb, device_idx
-
-        print(f"[DEBUG] vmfb not found at {self.falcon_vmfb_path.absolute()}")
-        if self.falcon_mlir_path.exists():
-            print(f"[DEBUG] mlir found at {self.falcon_mlir_path.absolute()}")
-            with open(self.falcon_mlir_path, "rb") as f:
-                bytecode = f.read()
-        else:
-            mlir_generated = False
-            print(
-                f"[DEBUG] mlir not found at {self.falcon_mlir_path.absolute()}"
-            )
-            if args.load_mlir_from_shark_tank:
-                # Downloading MLIR from shark_tank
-                print(f"[DEBUG] Trying to download mlir from shark_tank")
-                download_public_file(
-                    f"gs://shark_tank/falcon/sharded/falcon_{args.falcon_variant_to_use}/mlir/"
-                    + str(self.falcon_mlir_path),
-                    self.falcon_mlir_path.absolute(),
-                    single_file=True,
-                )
-                if self.falcon_mlir_path.exists():
-                    print(
-                        f"[DEBUG] mlir found at {self.falcon_mlir_path.absolute()}"
-                    )
-                    with open(self.falcon_mlir_path, "rb") as f:
-                        bytecode = f.read()
-                    mlir_generated = True
-
-            if not mlir_generated:
-                print(f"[DEBUG] generating MLIR locally")
-                if layer_id == "word_embeddings":
-                    f16_input_mask = [False]
-                elif layer_id in ["ln_f", "lm_head"]:
-                    f16_input_mask = [True]
-                elif "_" in layer_id or type(layer_id) == int:
-                    f16_input_mask = [True, True]
-                else:
-                    raise ValueError("Unsupported layer: ", layer_id)
-
-                print(f"[DEBUG] generating torchscript graph")
-                ts_graph = import_with_fx(
-                    layer,
-                    falconCompileInput,
-                    is_f16=True,
-                    f16_input_mask=f16_input_mask,
-                    mlir_type="torchscript",
-                    is_gptq=True,
-                )
-                del layer
-
-                print(f"[DEBUG] generating torch mlir")
-                module = torch_mlir.compile(
-                    ts_graph,
-                    falconCompileInput,
-                    torch_mlir.OutputType.LINALG_ON_TENSORS,
-                    use_tracing=False,
-                    verbose=False,
-                )
-                del ts_graph
-
-                print(f"[DEBUG] converting to bytecode")
-                bytecode_stream = BytesIO()
-                module.operation.write_bytecode(bytecode_stream)
-                bytecode = bytecode_stream.getvalue()
-                del module
-
-                f_ = open(self.falcon_mlir_path, "wb")
-                f_.write(bytecode)
-                print("Saved falcon mlir at ", str(self.falcon_mlir_path))
-                f_.close()
-                del bytecode
-
-        shark_module = SharkInference(
-            mlir_module=self.falcon_mlir_path,
-            device=self.device,
-            mlir_dialect="linalg",
-            device_idx=device_idx,
-        )
-        path = shark_module.save_module(
-            self.falcon_vmfb_path.parent.absolute(),
-            self.falcon_vmfb_path.stem,
-            extra_args=[
-                "--iree-vm-target-truncate-unsupported-floats",
-                "--iree-codegen-check-ir-before-llvm-conversion=false",
-                "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
-            ]
-            + [
-                "--iree-llvmcpu-use-fast-min-max-ops",
-            ]
-            if self.precision == "int4"
-            else [],
-            debug=self.debug,
-        )
-        print("Saved falcon vmfb at ", str(path))
-        shark_module.load_module(path)
-
-        return shark_module, device_idx
-
-    def compile(self):
-        sample_input_ids = torch.zeros([100], dtype=torch.int64)
-        sample_attention_mask = torch.zeros(
-            [1, 1, 100, 100], dtype=torch.float32
-        )
-        num_group_layers = int(
-            20 * (4 / args.num_shards)
-        )  # 4 is the number of default shards
-        sample_hidden_states = torch.zeros(
-            [1, 100, 14848], dtype=torch.float32
-        )
-
-        # Determine number of available devices
-        num_devices = 1
-        if self.device == "rocm":
-            import iree.runtime as ireert
-
-            haldriver = ireert.get_driver(self.device)
-            num_devices = len(haldriver.query_available_devices())
-            if num_devices < 2:
-                raise ValueError(
-                    "Cannot run Falcon-180B on a single ROCM device."
-                )
-
-        lm_head = LMHeadEmbeddingLayer(self.src_model.lm_head)
-        print("Compiling Layer lm_head")
-        shark_lm_head, _ = self.compile_layer(
-            lm_head,
-            [sample_hidden_states],
-            "lm_head",
-            device_idx=(0 % num_devices) % args.num_shards
-            if self.device == "rocm"
-            else None,
-        )
-        shark_lm_head = CompiledLMHeadEmbeddingLayer(shark_lm_head)
-
-        word_embedding = WordEmbeddingsLayer(
-            self.src_model.transformer.word_embeddings
-        )
-        print("Compiling Layer word_embeddings")
-        shark_word_embedding, _ = self.compile_layer(
-            word_embedding,
-            [sample_input_ids],
-            "word_embeddings",
-            device_idx=(1 % num_devices) % args.num_shards
-            if self.device == "rocm"
-            else None,
-        )
-        shark_word_embedding = CompiledWordEmbeddingsLayer(
-            shark_word_embedding
-        )
-
-        ln_f = LNFEmbeddingLayer(self.src_model.transformer.ln_f)
-        print("Compiling Layer ln_f")
-        shark_ln_f, _ = self.compile_layer(
-            ln_f,
-            [sample_hidden_states],
-            "ln_f",
-            device_idx=(2 % num_devices) % args.num_shards
-            if self.device == "rocm"
-            else None,
-        )
-        shark_ln_f = CompiledLNFEmbeddingLayer(shark_ln_f)
-
-        shark_layers = []
-        for i in range(
-            int(len(self.src_model.transformer.h) / num_group_layers)
-        ):
-            device_idx = i % num_devices if self.device == "rocm" else None
-            layer_id = i
-            layer_id = (
-                str(i * num_group_layers)
-                + "_"
-                + str((i + 1) * num_group_layers)
-            )
-            pytorch_class = FourWayShardingDecoderLayer
-            compiled_class = CompiledFourWayShardingDecoderLayer
-            if args.num_shards == 2:
-                pytorch_class = TwoWayShardingDecoderLayer
-                compiled_class = CompiledTwoWayShardingDecoderLayer
-
-            print("Compiling Layer {}".format(layer_id))
-            layer_i = self.src_model.transformer.h[
-                i * num_group_layers : (i + 1) * num_group_layers
-            ]
-
-            pytorch_layer_i = pytorch_class(
-                layer_i, args.falcon_variant_to_use
-            )
-            shark_module, device_idx = self.compile_layer(
-                pytorch_layer_i,
-                [sample_hidden_states, sample_attention_mask],
-                layer_id,
-                device_idx=device_idx,
-            )
-            shark_layer_i = compiled_class(
-                layer_id,
-                device_idx,
-                args.falcon_variant_to_use,
-                self.device,
-                self.precision,
-                shark_module,
-            )
-            shark_layers.append(shark_layer_i)
-
-        sharded_model = ShardedFalconModel(
-            self.src_model,
-            shark_layers,
-            shark_word_embedding,
-            shark_ln_f,
-            shark_lm_head,
-        )
-        return sharded_model
-
-    def generate(self, prompt):
-        model_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=self.max_padding_length,
-            add_special_tokens=False,
-            return_tensors="pt",
-        )
-        model_inputs["prompt_text"] = prompt
-
-        input_ids = model_inputs["input_ids"]
-        attention_mask = model_inputs.get("attention_mask", None)
-
-        # Allow empty prompts
-        if input_ids.shape[1] == 0:
-            input_ids = None
-            attention_mask = None
-
-        generate_kwargs = {
-            "max_length": self.max_num_tokens,
-            "do_sample": True,
-            "top_k": 10,
-            "num_return_sequences": 1,
-            "eos_token_id": 11,
-        }
-        generate_kwargs["input_ids"] = input_ids
-        generate_kwargs["attention_mask"] = attention_mask
-        generation_config_ = GenerationConfig.from_model_config(
-            self.src_model.config
-        )
-        generation_config = copy.deepcopy(generation_config_)
-        model_kwargs = generation_config.update(**generate_kwargs)
-
-        logits_processor = LogitsProcessorList()
-        stopping_criteria = StoppingCriteriaList()
-
-        eos_token_id = generation_config.eos_token_id
-        generation_config.pad_token_id = eos_token_id
-
-        (
-            inputs_tensor,
-            model_input_name,
-            model_kwargs,
-        ) = self.src_model._prepare_model_inputs(
-            None, generation_config.bos_token_id, model_kwargs
-        )
-
-        model_kwargs["output_attentions"] = generation_config.output_attentions
-        model_kwargs[
-            "output_hidden_states"
-        ] = generation_config.output_hidden_states
-        model_kwargs["use_cache"] = generation_config.use_cache
-
-        input_ids = (
-            inputs_tensor
-            if model_input_name == "input_ids"
-            else model_kwargs.pop("input_ids")
-        )
-
-        self.logits_processor = self.src_model._get_logits_processor(
-            generation_config=generation_config,
-            input_ids_seq_length=input_ids.shape[-1],
-            encoder_input_ids=inputs_tensor,
-            prefix_allowed_tokens_fn=None,
-            logits_processor=logits_processor,
-        )
-
-        self.stopping_criteria = self.src_model._get_stopping_criteria(
-            generation_config=generation_config,
-            stopping_criteria=stopping_criteria,
-        )
-
-        self.logits_warper = self.src_model._get_logits_warper(
-            generation_config
-        )
-
-        (
-            self.input_ids,
-            self.model_kwargs,
-        ) = self.src_model._expand_inputs_for_generation(
-            input_ids=input_ids,
-            expand_size=generation_config.num_return_sequences,  # 1
-            is_encoder_decoder=self.src_model.config.is_encoder_decoder,  # False
-            **model_kwargs,
-        )
-
-        if isinstance(eos_token_id, int):
-            eos_token_id = [eos_token_id]
-        self.eos_token_id_tensor = (
-            torch.tensor(eos_token_id) if eos_token_id is not None else None
-        )
-
-        self.pad_token_id = generation_config.pad_token_id
-        self.eos_token_id = eos_token_id
-
-        output_scores = generation_config.output_scores  # False
-        return_dict_in_generate = (
-            generation_config.return_dict_in_generate  # False
-        )
-
-        # init attention / hidden states / scores tuples
-        self.scores = (
-            () if (return_dict_in_generate and output_scores) else None
-        )
-
-        # keep track of which sequences are already finished
-        self.unfinished_sequences = torch.ones(
-            input_ids.shape[0], dtype=torch.long, device=input_ids.device
-        )
-
-        all_text = prompt
-
-        start = time.time()
-        count = 0
-        for i in range(self.max_num_tokens - 1):
-            count = count + 1
-
-            next_token = self.generate_new_token()
-            new_word = self.tokenizer.decode(
-                next_token.cpu().numpy(),
-                add_special_tokens=False,
-                skip_special_tokens=True,
-                clean_up_tokenization_spaces=True,
-            )
-
-            all_text = all_text + new_word
-
-            print(f"{new_word}", end="", flush=True)
-            print(f"{all_text}", end="", flush=True)
-
-            # if eos_token was found in one sentence, set sentence to finished
-            if self.eos_token_id_tensor is not None:
-                self.unfinished_sequences = self.unfinished_sequences.mul(
-                    next_token.tile(self.eos_token_id_tensor.shape[0], 1)
-                    .ne(self.eos_token_id_tensor.unsqueeze(1))
-                    .prod(dim=0)
-                )
-                # stop when each sentence is finished
-                if (
-                    self.unfinished_sequences.max() == 0
-                    or self.stopping_criteria(input_ids, self.scores)
-                ):
-                    break
-
-        end = time.time()
-        print(
-            "\n\nTime taken is {:.2f} seconds/token\n".format(
-                (end - start) / count
-            )
-        )
-
-        torch.cuda.empty_cache()
-        gc.collect()
-
-        return all_text
-
-    def generate_new_token(self):
-        model_inputs = self.src_model.prepare_inputs_for_generation(
-            self.input_ids, **self.model_kwargs
-        )
-        outputs = self.shark_model.forward(
-            input_ids=model_inputs["input_ids"],
-            attention_mask=model_inputs["attention_mask"],
-        )
-        if self.precision in ["fp16", "int4"]:
-            outputs = outputs.to(dtype=torch.float32)
-        next_token_logits = outputs
-
-        # pre-process distribution
-        next_token_scores = self.logits_processor(
-            self.input_ids, next_token_logits
-        )
-        next_token_scores = self.logits_warper(
-            self.input_ids, next_token_scores
-        )
-
-        # sample
-        probs = torch.nn.functional.softmax(next_token_scores, dim=-1)
-
-        next_token = torch.multinomial(probs, num_samples=1).squeeze(1)
-
-        # finished sentences should have their next token be a padding token
-        if self.eos_token_id is not None:
-            if self.pad_token_id is None:
-                raise ValueError(
-                    "If `eos_token_id` is defined, make sure that `pad_token_id` is defined."
-                )
-            next_token = (
-                next_token * self.unfinished_sequences
-                + self.pad_token_id * (1 - self.unfinished_sequences)
-            )
-
-        self.input_ids = torch.cat(
-            [self.input_ids, next_token[:, None]], dim=-1
-        )
-
-        self.model_kwargs["past_key_values"] = None
-        if "attention_mask" in self.model_kwargs:
-            attention_mask = self.model_kwargs["attention_mask"]
-            self.model_kwargs["attention_mask"] = torch.cat(
-                [
-                    attention_mask,
-                    attention_mask.new_ones((attention_mask.shape[0], 1)),
-                ],
-                dim=-1,
-            )
-
-        self.input_ids = self.input_ids[:, 1:]
-        self.model_kwargs["attention_mask"] = self.model_kwargs[
-            "attention_mask"
-        ][:, 1:]
-
-        return next_token
-
-
-class UnshardedFalcon(SharkLLMBase):
-    def __init__(
-        self,
-        model_name,
-        hf_model_path="tiiuae/falcon-7b-instruct",
-        hf_auth_token: str = "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",
-        max_num_tokens=150,
-        device="cuda",
-        precision="fp32",
-        falcon_mlir_path=None,
-        falcon_vmfb_path=None,
-        debug=False,
-    ) -> None:
-        super().__init__(model_name, hf_model_path, max_num_tokens)
-        print("hf_model_path: ", self.hf_model_path)
-
-        if "180b" in self.model_name and hf_auth_token == None:
-            raise ValueError(
-                """ HF auth token required for falcon-180b. Pass it using
-                --hf_auth_token flag. You can ask for the access to the model
-                here: https://huggingface.co/tiiuae/falcon-180B-chat."""
-            )
-        self.hf_auth_token = hf_auth_token
-        self.max_padding_length = 100
-        self.device = device
-        self.precision = precision
-        self.falcon_vmfb_path = falcon_vmfb_path
-        self.falcon_mlir_path = falcon_mlir_path
-        self.debug = debug
-        self.tokenizer = self.get_tokenizer()
-        self.src_model = self.get_src_model()
-        self.shark_model = self.compile()
-
-    def get_tokenizer(self):
-        tokenizer = AutoTokenizer.from_pretrained(
-            self.hf_model_path,
-            trust_remote_code=True,
-            token=self.hf_auth_token,
-        )
-        tokenizer.padding_side = "left"
-        tokenizer.pad_token_id = 11
-        return tokenizer
-
-    def get_src_model(self):
-        print("Loading src model: ", self.model_name)
-        kwargs = {
-            "torch_dtype": torch.float32,
-            "trust_remote_code": True,
-            "token": self.hf_auth_token,
-        }
-        if self.precision == "int4":
-            quantization_config = GPTQConfig(bits=4, disable_exllama=True)
-            kwargs["quantization_config"] = quantization_config
-            kwargs["device_map"] = "cpu"
-        falcon_model = AutoModelForCausalLM.from_pretrained(
-            self.hf_model_path, **kwargs
-        )
-        return falcon_model
-
-    def compile(self):
+    def compile_falcon(self):
         if args.use_precompiled_model:
             if not self.falcon_vmfb_path.exists():
                 # Downloading VMFB from shark_tank
@@ -706,37 +120,37 @@ class UnshardedFalcon(SharkLLMBase):
             if vmfb is not None:
                 return vmfb
 
-        print(f"[DEBUG] vmfb not found at {self.falcon_vmfb_path.absolute()}")
+        print(
+            f"[DEBUG] vmfb not found at {self.falcon_vmfb_path.absolute()}. Trying to work with"
+            f"[DEBUG] mlir path { self.falcon_mlir_path} {'exists' if self.falcon_mlir_path.exists() else 'does not exist'}"
+        )
         if self.falcon_mlir_path.exists():
-            print(f"[DEBUG] mlir found at {self.falcon_mlir_path.absolute()}")
             with open(self.falcon_mlir_path, "rb") as f:
                 bytecode = f.read()
         else:
             mlir_generated = False
-            print(
-                f"[DEBUG] mlir not found at {self.falcon_mlir_path.absolute()}"
+            # Downloading MLIR from shark_tank
+            download_public_file(
+                "gs://shark_tank/falcon/"
+                + "falcon_"
+                + args.falcon_variant_to_use
+                + "_"
+                + self.precision
+                + ".mlir",
+                self.falcon_mlir_path.absolute(),
+                single_file=True,
             )
-            if args.load_mlir_from_shark_tank:
-                # Downloading MLIR from shark_tank
-                print(f"[DEBUG] Trying to download mlir from shark_tank")
-                download_public_file(
-                    "gs://shark_tank/falcon/"
-                    + "falcon_"
-                    + args.falcon_variant_to_use
-                    + "_"
-                    + self.precision
-                    + ".mlir",
-                    self.falcon_mlir_path.absolute(),
-                    single_file=True,
+            if self.falcon_mlir_path.exists():
+                with open(self.falcon_mlir_path, "rb") as f:
+                    bytecode = f.read()
+                mlir_generated = True
+            else:
+                raise ValueError(
+                    f"MLIR not found at {self.falcon_mlir_path.absolute()}"
+                    " after downloading! Please check path and try again"
                 )
-                if self.falcon_mlir_path.exists():
-                    print(
-                        f"[DEBUG] mlir found at {self.falcon_mlir_path.absolute()}"
-                    )
-                    mlir_generated = True
 
             if not mlir_generated:
-                print(f"[DEBUG] generating MLIR locally")
                 compilation_input_ids = torch.randint(
                     low=1, high=10000, size=(1, 100)
                 )
@@ -753,10 +167,9 @@ class UnshardedFalcon(SharkLLMBase):
                 ts_graph = import_with_fx(
                     model,
                     falconCompileInput,
-                    is_f16=self.precision in ["fp16", "int4"],
+                    is_f16=self.precision == "fp16",
                     f16_input_mask=[False, False],
                     mlir_type="torchscript",
-                    is_gptq=self.precision == "int4",
                 )
                 del model
                 print(f"[DEBUG] generating torch mlir")
@@ -776,37 +189,35 @@ class UnshardedFalcon(SharkLLMBase):
                 bytecode = bytecode_stream.getvalue()
                 del module
 
-                f_ = open(self.falcon_mlir_path, "wb")
-                f_.write(bytecode)
-                print("Saved falcon mlir at ", str(self.falcon_mlir_path))
+                print(f"[DEBUG] writing mlir to file")
+                with open(f"{self.model_name}.mlir", "wb") as f_:
+                    with redirect_stdout(f_):
+                        print(module.operation.get_asm())
                 f_.close()
-                del bytecode
 
         shark_module = SharkInference(
-            mlir_module=self.falcon_mlir_path,
-            device=self.device,
-            mlir_dialect="linalg",
+            mlir_module=bytecode, device=self.device, mlir_dialect="linalg"
         )
         path = shark_module.save_module(
             self.falcon_vmfb_path.parent.absolute(),
             self.falcon_vmfb_path.stem,
             extra_args=[
+                "--iree-hal-dump-executable-sources-to=ies",
                 "--iree-vm-target-truncate-unsupported-floats",
                 "--iree-codegen-check-ir-before-llvm-conversion=false",
                 "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
-            ]
-            + [
-                "--iree-llvmcpu-use-fast-min-max-ops",
-            ]
-            if self.precision == "int4"
-            else [],
-            debug=self.debug,
+                "--iree-spirv-index-bits=64",
+            ],
         )
         print("Saved falcon vmfb at ", str(path))
         shark_module.load_module(path)
 
         return shark_module
 
+    def compile(self):
+        falcon_shark_model = self.compile_falcon()
+        return falcon_shark_model
+
     def generate(self, prompt):
         model_inputs = self.tokenizer(
             prompt,
@@ -934,11 +345,7 @@ class UnshardedFalcon(SharkLLMBase):
 
         all_text = prompt
 
-        start = time.time()
-        count = 0
         for i in range(self.max_num_tokens - 1):
-            count = count + 1
-
             next_token = self.generate_new_token()
             new_word = self.tokenizer.decode(
                 next_token.cpu().numpy(),
@@ -965,13 +372,6 @@ class UnshardedFalcon(SharkLLMBase):
                 ):
                     break
 
-        end = time.time()
-        print(
-            "\n\nTime taken is {:.2f} seconds/token\n".format(
-                (end - start) / count
-            )
-        )
-
         torch.cuda.empty_cache()
         gc.collect()
 
@@ -987,7 +387,7 @@ class UnshardedFalcon(SharkLLMBase):
                 (model_inputs["input_ids"], model_inputs["attention_mask"]),
             )
         )
-        if self.precision in ["fp16", "int4"]:
+        if self.precision == "fp16":
             outputs = outputs.to(dtype=torch.float32)
         next_token_logits = outputs
 
@@ -1066,39 +466,18 @@ if __name__ == "__main__":
         else Path(args.falcon_vmfb_path)
     )
 
-    if args.precision == "int4":
-        if args.falcon_variant_to_use == "180b":
-            hf_model_path_value = "TheBloke/Falcon-180B-Chat-GPTQ"
-        else:
-            hf_model_path_value = (
-                "TheBloke/falcon-"
-                + args.falcon_variant_to_use
-                + "-instruct-GPTQ"
-            )
-    else:
-        if args.falcon_variant_to_use == "180b":
-            hf_model_path_value = "tiiuae/falcon-180B-chat"
-        else:
-            hf_model_path_value = (
-                "tiiuae/falcon-" + args.falcon_variant_to_use + "-instruct"
-            )
+    falcon = Falcon(
+        "falcon_" + args.falcon_variant_to_use,
+        hf_model_path="tiiuae/falcon-"
+        + args.falcon_variant_to_use
+        + "-instruct",
+        device=args.device,
+        precision=args.precision,
+        falcon_mlir_path=falcon_mlir_path,
+        falcon_vmfb_path=falcon_vmfb_path,
+    )
 
-    if not args.sharded:
-        falcon = UnshardedFalcon(
-            model_name="falcon_" + args.falcon_variant_to_use,
-            hf_model_path=hf_model_path_value,
-            device=args.device,
-            precision=args.precision,
-            falcon_mlir_path=falcon_mlir_path,
-            falcon_vmfb_path=falcon_vmfb_path,
-        )
-    else:
-        falcon = ShardedFalcon(
-            model_name="falcon_" + args.falcon_variant_to_use,
-            hf_model_path=hf_model_path_value,
-            device=args.device,
-            precision=args.precision,
-        )
+    import gc
 
     default_prompt_text = "Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Giraftron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\nDaniel: Hello, Girafatron!\nGirafatron:"
     continue_execution = True
@@ -1118,11 +497,7 @@ if __name__ == "__main__":
             prompt = input("Please enter the prompt text: ")
         print("\nPrompt Text: ", prompt)
 
-        prompt_template = f"""A helpful assistant who helps the user with any questions asked.
-        User: {prompt}
-        Assistant:"""
-
-        res_str = falcon.generate(prompt_template)
+        res_str = falcon.generate(prompt)
         torch.cuda.empty_cache()
         gc.collect()
         print(

apps/language_models/src/pipelines/minigpt4_utils/blip_processors.py

@@ -1,68 +0,0 @@
-"""
- Copyright (c) 2022, salesforce.com, inc.
- All rights reserved.
- SPDX-License-Identifier: BSD-3-Clause
- For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause
-"""
-from omegaconf import OmegaConf
-from torchvision import transforms
-from torchvision.transforms.functional import InterpolationMode
-
-
-class BaseProcessor:
-    def __init__(self):
-        self.transform = lambda x: x
-        return
-
-    def __call__(self, item):
-        return self.transform(item)
-
-    @classmethod
-    def from_config(cls, cfg=None):
-        return cls()
-
-    def build(self, **kwargs):
-        cfg = OmegaConf.create(kwargs)
-
-        return self.from_config(cfg)
-
-
-class BlipImageBaseProcessor(BaseProcessor):
-    def __init__(self, mean=None, std=None):
-        if mean is None:
-            mean = (0.48145466, 0.4578275, 0.40821073)
-        if std is None:
-            std = (0.26862954, 0.26130258, 0.27577711)
-
-        self.normalize = transforms.Normalize(mean, std)
-
-
-class Blip2ImageEvalProcessor(BlipImageBaseProcessor):
-    def __init__(self, image_size=224, mean=None, std=None):
-        super().__init__(mean=mean, std=std)
-
-        self.transform = transforms.Compose(
-            [
-                transforms.Resize(
-                    (image_size, image_size),
-                    interpolation=InterpolationMode.BICUBIC,
-                ),
-                transforms.ToTensor(),
-                self.normalize,
-            ]
-        )
-
-    def __call__(self, item):
-        return self.transform(item)
-
-    @classmethod
-    def from_config(cls, cfg=None):
-        if cfg is None:
-            cfg = OmegaConf.create()
-
-        image_size = cfg.get("image_size", 224)
-
-        mean = cfg.get("mean", None)
-        std = cfg.get("std", None)
-
-        return cls(image_size=image_size, mean=mean, std=std)

apps/language_models/src/pipelines/minigpt4_utils/configs/cc_sbu_align.yaml

@@ -1,5 +0,0 @@
-datasets:
-  cc_sbu_align:
-    data_type: images
-    build_info:
-      storage: /path/to/cc_sbu_align/

apps/language_models/src/pipelines/minigpt4_utils/configs/minigpt4.yaml

@@ -1,33 +0,0 @@
-model:
-  arch: mini_gpt4
-
-  # vit encoder
-  image_size: 224
-  drop_path_rate: 0
-  use_grad_checkpoint: False
-  vit_precision: "fp16"
-  freeze_vit: True
-  freeze_qformer: True
-
-  # Q-Former
-  num_query_token: 32
-
-  # Vicuna
-  llama_model: "lmsys/vicuna-7b-v1.3"
-
-  # generation configs
-  prompt: ""
-
-preprocess:
-    vis_processor:
-        train:
-          name: "blip2_image_train"
-          image_size: 224
-        eval:
-          name: "blip2_image_eval"
-          image_size: 224
-    text_processor:
-        train:
-          name: "blip_caption"
-        eval:
-          name: "blip_caption"

apps/language_models/src/pipelines/minigpt4_utils/configs/minigpt4_eval.yaml

@@ -1,25 +0,0 @@
-model:
-  arch: mini_gpt4
-  model_type: pretrain_vicuna
-  freeze_vit: True
-  freeze_qformer: True
-  max_txt_len: 160
-  end_sym: "###"
-  low_resource: False
-  prompt_path: "apps/language_models/src/pipelines/minigpt4_utils/prompts/alignment.txt"
-  prompt_template: '###Human: {} ###Assistant: '
-  ckpt: 'prerained_minigpt4_7b.pth'
-
-
-datasets:
-  cc_sbu_align:
-    vis_processor:
-      train:
-        name: "blip2_image_eval"
-        image_size: 224
-    text_processor:
-      train:
-        name: "blip_caption"
-
-run:
-  task: image_text_pretrain

apps/language_models/src/pipelines/minigpt4_utils/eva_vit.py

@@ -1,629 +0,0 @@
-# Based on EVA, BEIT, timm and DeiT code bases
-# https://github.com/baaivision/EVA
-# https://github.com/rwightman/pytorch-image-models/tree/master/timm
-# https://github.com/microsoft/unilm/tree/master/beit
-# https://github.com/facebookresearch/deit/
-# https://github.com/facebookresearch/dino
-# --------------------------------------------------------'
-import math
-import requests
-from functools import partial
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint as checkpoint
-from timm.models.layers import drop_path, to_2tuple, trunc_normal_
-
-
-def _cfg(url="", **kwargs):
-    return {
-        "url": url,
-        "num_classes": 1000,
-        "input_size": (3, 224, 224),
-        "pool_size": None,
-        "crop_pct": 0.9,
-        "interpolation": "bicubic",
-        "mean": (0.5, 0.5, 0.5),
-        "std": (0.5, 0.5, 0.5),
-        **kwargs,
-    }
-
-
-class DropPath(nn.Module):
-    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
-
-    def __init__(self, drop_prob=None):
-        super(DropPath, self).__init__()
-        self.drop_prob = drop_prob
-
-    def forward(self, x):
-        return drop_path(x, self.drop_prob, self.training)
-
-    def extra_repr(self) -> str:
-        return "p={}".format(self.drop_prob)
-
-
-class Mlp(nn.Module):
-    def __init__(
-        self,
-        in_features,
-        hidden_features=None,
-        out_features=None,
-        act_layer=nn.GELU,
-        drop=0.0,
-    ):
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features)
-        self.act = act_layer()
-        self.fc2 = nn.Linear(hidden_features, out_features)
-        self.drop = nn.Dropout(drop)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        # x = self.drop(x)
-        # commit this for the orignal BERT implement
-        x = self.fc2(x)
-        x = self.drop(x)
-        return x
-
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        dim,
-        num_heads=8,
-        qkv_bias=False,
-        qk_scale=None,
-        attn_drop=0.0,
-        proj_drop=0.0,
-        window_size=None,
-        attn_head_dim=None,
-    ):
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        if attn_head_dim is not None:
-            head_dim = attn_head_dim
-        all_head_dim = head_dim * self.num_heads
-        self.scale = qk_scale or head_dim**-0.5
-
-        self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
-        if qkv_bias:
-            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
-            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
-        else:
-            self.q_bias = None
-            self.v_bias = None
-
-        if window_size:
-            self.window_size = window_size
-            self.num_relative_distance = (2 * window_size[0] - 1) * (
-                2 * window_size[1] - 1
-            ) + 3
-            self.relative_position_bias_table = nn.Parameter(
-                torch.zeros(self.num_relative_distance, num_heads)
-            )  # 2*Wh-1 * 2*Ww-1, nH
-            # cls to token & token 2 cls & cls to cls
-
-            # get pair-wise relative position index for each token inside the window
-            coords_h = torch.arange(window_size[0])
-            coords_w = torch.arange(window_size[1])
-            coords = torch.stack(
-                torch.meshgrid([coords_h, coords_w])
-            )  # 2, Wh, Ww
-            coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
-            relative_coords = (
-                coords_flatten[:, :, None] - coords_flatten[:, None, :]
-            )  # 2, Wh*Ww, Wh*Ww
-            relative_coords = relative_coords.permute(
-                1, 2, 0
-            ).contiguous()  # Wh*Ww, Wh*Ww, 2
-            relative_coords[:, :, 0] += (
-                window_size[0] - 1
-            )  # shift to start from 0
-            relative_coords[:, :, 1] += window_size[1] - 1
-            relative_coords[:, :, 0] *= 2 * window_size[1] - 1
-            relative_position_index = torch.zeros(
-                size=(window_size[0] * window_size[1] + 1,) * 2,
-                dtype=relative_coords.dtype,
-            )
-            relative_position_index[1:, 1:] = relative_coords.sum(
-                -1
-            )  # Wh*Ww, Wh*Ww
-            relative_position_index[0, 0:] = self.num_relative_distance - 3
-            relative_position_index[0:, 0] = self.num_relative_distance - 2
-            relative_position_index[0, 0] = self.num_relative_distance - 1
-
-            self.register_buffer(
-                "relative_position_index", relative_position_index
-            )
-        else:
-            self.window_size = None
-            self.relative_position_bias_table = None
-            self.relative_position_index = None
-
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.proj = nn.Linear(all_head_dim, dim)
-        self.proj_drop = nn.Dropout(proj_drop)
-
-    def forward(self, x, rel_pos_bias=None):
-        B, N, C = x.shape
-        qkv_bias = None
-        if self.q_bias is not None:
-            qkv_bias = torch.cat(
-                (
-                    self.q_bias,
-                    torch.zeros_like(self.v_bias, requires_grad=False),
-                    self.v_bias,
-                )
-            )
-        # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
-        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
-        q, k, v = (
-            qkv[0],
-            qkv[1],
-            qkv[2],
-        )  # make torchscript happy (cannot use tensor as tuple)
-
-        q = q * self.scale
-        attn = q @ k.transpose(-2, -1)
-
-        if self.relative_position_bias_table is not None:
-            relative_position_bias = self.relative_position_bias_table[
-                self.relative_position_index.view(-1)
-            ].view(
-                self.window_size[0] * self.window_size[1] + 1,
-                self.window_size[0] * self.window_size[1] + 1,
-                -1,
-            )  # Wh*Ww,Wh*Ww,nH
-            relative_position_bias = relative_position_bias.permute(
-                2, 0, 1
-            ).contiguous()  # nH, Wh*Ww, Wh*Ww
-            attn = attn + relative_position_bias.unsqueeze(0)
-
-        if rel_pos_bias is not None:
-            attn = attn + rel_pos_bias
-
-        attn = attn.softmax(dim=-1)
-        attn = self.attn_drop(attn)
-
-        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-class Block(nn.Module):
-    def __init__(
-        self,
-        dim,
-        num_heads,
-        mlp_ratio=4.0,
-        qkv_bias=False,
-        qk_scale=None,
-        drop=0.0,
-        attn_drop=0.0,
-        drop_path=0.0,
-        init_values=None,
-        act_layer=nn.GELU,
-        norm_layer=nn.LayerNorm,
-        window_size=None,
-        attn_head_dim=None,
-    ):
-        super().__init__()
-        self.norm1 = norm_layer(dim)
-        self.attn = Attention(
-            dim,
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            qk_scale=qk_scale,
-            attn_drop=attn_drop,
-            proj_drop=drop,
-            window_size=window_size,
-            attn_head_dim=attn_head_dim,
-        )
-        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
-        self.drop_path = (
-            DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
-        )
-        self.norm2 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = Mlp(
-            in_features=dim,
-            hidden_features=mlp_hidden_dim,
-            act_layer=act_layer,
-            drop=drop,
-        )
-
-        if init_values is not None and init_values > 0:
-            self.gamma_1 = nn.Parameter(
-                init_values * torch.ones((dim)), requires_grad=True
-            )
-            self.gamma_2 = nn.Parameter(
-                init_values * torch.ones((dim)), requires_grad=True
-            )
-        else:
-            self.gamma_1, self.gamma_2 = None, None
-
-    def forward(self, x, rel_pos_bias=None):
-        if self.gamma_1 is None:
-            x = x + self.drop_path(
-                self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias)
-            )
-            x = x + self.drop_path(self.mlp(self.norm2(x)))
-        else:
-            x = x + self.drop_path(
-                self.gamma_1
-                * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias)
-            )
-            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
-        return x
-
-
-class PatchEmbed(nn.Module):
-    """Image to Patch Embedding"""
-
-    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
-        super().__init__()
-        img_size = to_2tuple(img_size)
-        patch_size = to_2tuple(patch_size)
-        num_patches = (img_size[1] // patch_size[1]) * (
-            img_size[0] // patch_size[0]
-        )
-        self.patch_shape = (
-            img_size[0] // patch_size[0],
-            img_size[1] // patch_size[1],
-        )
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.num_patches = num_patches
-
-        self.proj = nn.Conv2d(
-            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
-        )
-
-    def forward(self, x, **kwargs):
-        B, C, H, W = x.shape
-        # FIXME look at relaxing size constraints
-        assert (
-            H == self.img_size[0] and W == self.img_size[1]
-        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
-        x = self.proj(x).flatten(2).transpose(1, 2)
-        return x
-
-
-class RelativePositionBias(nn.Module):
-    def __init__(self, window_size, num_heads):
-        super().__init__()
-        self.window_size = window_size
-        self.num_relative_distance = (2 * window_size[0] - 1) * (
-            2 * window_size[1] - 1
-        ) + 3
-        self.relative_position_bias_table = nn.Parameter(
-            torch.zeros(self.num_relative_distance, num_heads)
-        )  # 2*Wh-1 * 2*Ww-1, nH
-        # cls to token & token 2 cls & cls to cls
-
-        # get pair-wise relative position index for each token inside the window
-        coords_h = torch.arange(window_size[0])
-        coords_w = torch.arange(window_size[1])
-        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
-        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
-        relative_coords = (
-            coords_flatten[:, :, None] - coords_flatten[:, None, :]
-        )  # 2, Wh*Ww, Wh*Ww
-        relative_coords = relative_coords.permute(
-            1, 2, 0
-        ).contiguous()  # Wh*Ww, Wh*Ww, 2
-        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
-        relative_coords[:, :, 1] += window_size[1] - 1
-        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
-        relative_position_index = torch.zeros(
-            size=(window_size[0] * window_size[1] + 1,) * 2,
-            dtype=relative_coords.dtype,
-        )
-        relative_position_index[1:, 1:] = relative_coords.sum(
-            -1
-        )  # Wh*Ww, Wh*Ww
-        relative_position_index[0, 0:] = self.num_relative_distance - 3
-        relative_position_index[0:, 0] = self.num_relative_distance - 2
-        relative_position_index[0, 0] = self.num_relative_distance - 1
-
-        self.register_buffer(
-            "relative_position_index", relative_position_index
-        )
-
-        # trunc_normal_(self.relative_position_bias_table, std=.02)
-
-    def forward(self):
-        relative_position_bias = self.relative_position_bias_table[
-            self.relative_position_index.view(-1)
-        ].view(
-            self.window_size[0] * self.window_size[1] + 1,
-            self.window_size[0] * self.window_size[1] + 1,
-            -1,
-        )  # Wh*Ww,Wh*Ww,nH
-        return relative_position_bias.permute(
-            2, 0, 1
-        ).contiguous()  # nH, Wh*Ww, Wh*Ww
-
-
-class VisionTransformer(nn.Module):
-    """Vision Transformer with support for patch or hybrid CNN input stage"""
-
-    def __init__(
-        self,
-        img_size=224,
-        patch_size=16,
-        in_chans=3,
-        num_classes=1000,
-        embed_dim=768,
-        depth=12,
-        num_heads=12,
-        mlp_ratio=4.0,
-        qkv_bias=False,
-        qk_scale=None,
-        drop_rate=0.0,
-        attn_drop_rate=0.0,
-        drop_path_rate=0.0,
-        norm_layer=nn.LayerNorm,
-        init_values=None,
-        use_abs_pos_emb=True,
-        use_rel_pos_bias=False,
-        use_shared_rel_pos_bias=False,
-        use_mean_pooling=True,
-        init_scale=0.001,
-        use_checkpoint=False,
-    ):
-        super().__init__()
-        self.image_size = img_size
-        self.num_classes = num_classes
-        self.num_features = (
-            self.embed_dim
-        ) = embed_dim  # num_features for consistency with other models
-
-        self.patch_embed = PatchEmbed(
-            img_size=img_size,
-            patch_size=patch_size,
-            in_chans=in_chans,
-            embed_dim=embed_dim,
-        )
-        num_patches = self.patch_embed.num_patches
-
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        if use_abs_pos_emb:
-            self.pos_embed = nn.Parameter(
-                torch.zeros(1, num_patches + 1, embed_dim)
-            )
-        else:
-            self.pos_embed = None
-        self.pos_drop = nn.Dropout(p=drop_rate)
-
-        if use_shared_rel_pos_bias:
-            self.rel_pos_bias = RelativePositionBias(
-                window_size=self.patch_embed.patch_shape, num_heads=num_heads
-            )
-        else:
-            self.rel_pos_bias = None
-        self.use_checkpoint = use_checkpoint
-
-        dpr = [
-            x.item() for x in torch.linspace(0, drop_path_rate, depth)
-        ]  # stochastic depth decay rule
-        self.use_rel_pos_bias = use_rel_pos_bias
-        self.blocks = nn.ModuleList(
-            [
-                Block(
-                    dim=embed_dim,
-                    num_heads=num_heads,
-                    mlp_ratio=mlp_ratio,
-                    qkv_bias=qkv_bias,
-                    qk_scale=qk_scale,
-                    drop=drop_rate,
-                    attn_drop=attn_drop_rate,
-                    drop_path=dpr[i],
-                    norm_layer=norm_layer,
-                    init_values=init_values,
-                    window_size=self.patch_embed.patch_shape
-                    if use_rel_pos_bias
-                    else None,
-                )
-                for i in range(depth)
-            ]
-        )
-        #         self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
-        #         self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
-        #         self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
-
-        if self.pos_embed is not None:
-            trunc_normal_(self.pos_embed, std=0.02)
-        trunc_normal_(self.cls_token, std=0.02)
-        # trunc_normal_(self.mask_token, std=.02)
-        #         if isinstance(self.head, nn.Linear):
-        #             trunc_normal_(self.head.weight, std=.02)
-        self.apply(self._init_weights)
-        self.fix_init_weight()
-
-    #         if isinstance(self.head, nn.Linear):
-    #             self.head.weight.data.mul_(init_scale)
-    #             self.head.bias.data.mul_(init_scale)
-
-    def fix_init_weight(self):
-        def rescale(param, layer_id):
-            param.div_(math.sqrt(2.0 * layer_id))
-
-        for layer_id, layer in enumerate(self.blocks):
-            rescale(layer.attn.proj.weight.data, layer_id + 1)
-            rescale(layer.mlp.fc2.weight.data, layer_id + 1)
-
-    def _init_weights(self, m):
-        if isinstance(m, nn.Linear):
-            trunc_normal_(m.weight, std=0.02)
-            if isinstance(m, nn.Linear) and m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.LayerNorm):
-            nn.init.constant_(m.bias, 0)
-            nn.init.constant_(m.weight, 1.0)
-
-    def get_classifier(self):
-        return self.head
-
-    def reset_classifier(self, num_classes, global_pool=""):
-        self.num_classes = num_classes
-        self.head = (
-            nn.Linear(self.embed_dim, num_classes)
-            if num_classes > 0
-            else nn.Identity()
-        )
-
-    def forward_features(self, x):
-        x = self.patch_embed(x)
-        batch_size, seq_len, _ = x.size()
-
-        cls_tokens = self.cls_token.expand(
-            batch_size, -1, -1
-        )  # stole cls_tokens impl from Phil Wang, thanks
-        x = torch.cat((cls_tokens, x), dim=1)
-        if self.pos_embed is not None:
-            x = x + self.pos_embed
-        x = self.pos_drop(x)
-
-        rel_pos_bias = (
-            self.rel_pos_bias() if self.rel_pos_bias is not None else None
-        )
-        for blk in self.blocks:
-            if self.use_checkpoint:
-                x = checkpoint.checkpoint(blk, x, rel_pos_bias)
-            else:
-                x = blk(x, rel_pos_bias)
-        return x
-
-    #         x = self.norm(x)
-
-    #         if self.fc_norm is not None:
-    #             t = x[:, 1:, :]
-    #             return self.fc_norm(t.mean(1))
-    #         else:
-    #             return x[:, 0]
-
-    def forward(self, x):
-        x = self.forward_features(x)
-        #         x = self.head(x)
-        return x
-
-    def get_intermediate_layers(self, x):
-        x = self.patch_embed(x)
-        batch_size, seq_len, _ = x.size()
-
-        cls_tokens = self.cls_token.expand(
-            batch_size, -1, -1
-        )  # stole cls_tokens impl from Phil Wang, thanks
-        x = torch.cat((cls_tokens, x), dim=1)
-        if self.pos_embed is not None:
-            x = x + self.pos_embed
-        x = self.pos_drop(x)
-
-        features = []
-        rel_pos_bias = (
-            self.rel_pos_bias() if self.rel_pos_bias is not None else None
-        )
-        for blk in self.blocks:
-            x = blk(x, rel_pos_bias)
-            features.append(x)
-
-        return features
-
-
-def interpolate_pos_embed(model, checkpoint_model):
-    if "pos_embed" in checkpoint_model:
-        pos_embed_checkpoint = checkpoint_model["pos_embed"].float()
-        embedding_size = pos_embed_checkpoint.shape[-1]
-        num_patches = model.patch_embed.num_patches
-        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
-        # height (== width) for the checkpoint position embedding
-        orig_size = int(
-            (pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5
-        )
-        # height (== width) for the new position embedding
-        new_size = int(num_patches**0.5)
-        # class_token and dist_token are kept unchanged
-        if orig_size != new_size:
-            print(
-                "Position interpolate from %dx%d to %dx%d"
-                % (orig_size, orig_size, new_size, new_size)
-            )
-            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
-            # only the position tokens are interpolated
-            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
-            pos_tokens = pos_tokens.reshape(
-                -1, orig_size, orig_size, embedding_size
-            ).permute(0, 3, 1, 2)
-            pos_tokens = torch.nn.functional.interpolate(
-                pos_tokens,
-                size=(new_size, new_size),
-                mode="bicubic",
-                align_corners=False,
-            )
-            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
-            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
-            checkpoint_model["pos_embed"] = new_pos_embed
-
-
-def convert_weights_to_fp16(model: nn.Module):
-    """Convert applicable model parameters to fp16"""
-
-    def _convert_weights_to_fp16(l):
-        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
-            # l.weight.data = l.weight.data.half()
-            l.weight.data = l.weight.data
-            if l.bias is not None:
-                # l.bias.data = l.bias.data.half()
-                l.bias.data = l.bias.data
-
-    #         if isinstance(l, (nn.MultiheadAttention, Attention)):
-    #             for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
-    #                 tensor = getattr(l, attr)
-    #                 if tensor is not None:
-    #                     tensor.data = tensor.data.half()
-
-    model.apply(_convert_weights_to_fp16)
-
-
-def create_eva_vit_g(
-    img_size=224, drop_path_rate=0.4, use_checkpoint=False, precision="fp16"
-):
-    model = VisionTransformer(
-        img_size=img_size,
-        patch_size=14,
-        use_mean_pooling=False,
-        embed_dim=1408,
-        depth=39,
-        num_heads=1408 // 88,
-        mlp_ratio=4.3637,
-        qkv_bias=True,
-        drop_path_rate=drop_path_rate,
-        norm_layer=partial(nn.LayerNorm, eps=1e-6),
-        use_checkpoint=use_checkpoint,
-    )
-    url = "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/eva_vit_g.pth"
-
-    local_filename = "eva_vit_g.pth"
-    response = requests.get(url)
-    if response.status_code == 200:
-        with open(local_filename, "wb") as f:
-            f.write(response.content)
-        print("File downloaded successfully.")
-    state_dict = torch.load(local_filename, map_location="cpu")
-    interpolate_pos_embed(model, state_dict)
-
-    incompatible_keys = model.load_state_dict(state_dict, strict=False)
-
-    if precision == "fp16":
-        #         model.to("cuda")
-        convert_weights_to_fp16(model)
-    return model

apps/language_models/src/pipelines/minigpt4_utils/prompts/alignment.txt

@@ -1,4 +0,0 @@
-<Img><ImageHere></Img> Describe this image in detail.
-<Img><ImageHere></Img> Take a look at this image and describe what you notice.
-<Img><ImageHere></Img> Please provide a detailed description of the picture.
-<Img><ImageHere></Img> Could you describe the contents of this image for me?

apps/language_models/src/pipelines/stablelm_pipeline.py

@@ -32,13 +32,11 @@ class SharkStableLM(SharkLLMBase):
         max_num_tokens=512,
         device="cuda",
         precision="fp32",
-        debug="False",
     ) -> None:
         super().__init__(model_name, hf_model_path, max_num_tokens)
         self.max_sequence_len = 256
         self.device = device
         self.precision = precision
-        self.debug = debug
         self.tokenizer = self.get_tokenizer()
         self.shark_model = self.compile()
 
@@ -113,7 +111,7 @@ class SharkStableLM(SharkLLMBase):
         shark_module.compile()
 
         path = shark_module.save_module(
-            vmfb_path.parent.absolute(), vmfb_path.stem, debug=self.debug
+            vmfb_path.parent.absolute(), vmfb_path.stem
         )
         print("Saved vmfb at ", str(path))
 

apps/language_models/src/pipelines/vicuna_pipeline.py

@@ -0,0 +1,559 @@
+from apps.language_models.src.model_wrappers.vicuna_model import (
+    FirstVicuna,
+    SecondVicuna,
+)
+from apps.language_models.src.pipelines.SharkLLMBase import SharkLLMBase
+from apps.language_models.utils import (
+    get_vmfb_from_path,
+)
+
+from io import BytesIO
+from pathlib import Path
+from shark.shark_downloader import download_public_file
+from shark.shark_importer import import_with_fx
+from shark.shark_inference import SharkInference
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+import re
+import torch
+import torch_mlir
+import os
+
+
+class Vicuna(SharkLLMBase):
+    def __init__(
+        self,
+        model_name,
+        hf_model_path="TheBloke/vicuna-7B-1.1-HF",
+        max_num_tokens=512,
+        device="cuda",
+        precision="fp32",
+        first_vicuna_mlir_path=Path("first_vicuna.mlir"),
+        second_vicuna_mlir_path=Path("second_vicuna.mlir"),
+        first_vicuna_vmfb_path=Path("first_vicuna.vmfb"),
+        second_vicuna_vmfb_path=Path("second_vicuna.vmfb"),
+        load_mlir_from_shark_tank=True,
+    ) -> None:
+        super().__init__(model_name, hf_model_path, max_num_tokens)
+        self.max_sequence_length = 256
+        self.device = device
+        self.precision = precision
+        self.first_vicuna_vmfb_path = first_vicuna_vmfb_path
+        self.second_vicuna_vmfb_path = second_vicuna_vmfb_path
+        self.first_vicuna_mlir_path = first_vicuna_mlir_path
+        self.second_vicuna_mlir_path = second_vicuna_mlir_path
+        self.tokenizer = self.get_tokenizer()
+        self.shark_model = self.compile()
+        self.load_mlir_from_shark_tank = load_mlir_from_shark_tank
+
+    def get_tokenizer(self):
+        tokenizer = AutoTokenizer.from_pretrained(
+            self.hf_model_path, use_fast=False
+        )
+        return tokenizer
+
+    def get_src_model(self):
+        kwargs = {"torch_dtype": torch.float}
+        vicuna_model = AutoModelForCausalLM.from_pretrained(
+            self.hf_model_path, **kwargs
+        )
+        return vicuna_model
+
+    def compile_first_vicuna(self):
+        vmfb = get_vmfb_from_path(
+            self.first_vicuna_vmfb_path, self.device, "tm_tensor"
+        )
+        if vmfb is not None:
+            return vmfb
+
+        # Compilation path needs some more work before it is functional
+
+        print(
+            f"[DEBUG] vmfb not found at {self.first_vicuna_vmfb_path.absolute()}. Trying to work with"
+            f"[DEBUG] mlir path { self.first_vicuna_mlir_path} {'exists' if self.first_vicuna_mlir_path.exists() else 'does not exist'}"
+        )
+        if self.first_vicuna_mlir_path.exists():
+            with open(self.first_vicuna_mlir_path, "rb") as f:
+                bytecode = f.read()
+        else:
+            mlir_generated = False
+            if self.load_mlir_from_shark_tank:
+                if self.precision == "fp32":
+                    # download MLIR from shark_tank for fp32
+                    download_public_file(
+                        "gs://shark_tank/vicuna/unsharded/mlir/first_vicuna.mlir",
+                        self.first_vicuna_mlir_path.absolute(),
+                        single_file=True,
+                    )
+                    if self.first_vicuna_mlir_path.exists():
+                        with open(self.first_vicuna_mlir_path, "rb") as f:
+                            bytecode = f.read()
+                        mlir_generated = True
+                    else:
+                        raise ValueError(
+                            f"MLIR not found at {self.first_vicuna_mlir_path.absolute()}"
+                            " after downloading! Please check path and try again"
+                        )
+                else:
+                    print(
+                        "Only fp32 mlir added to tank, generating mlir on device."
+                    )
+
+            if not mlir_generated:
+                compilation_prompt = "".join(["0" for _ in range(17)])
+                compilation_input_ids = self.tokenizer(
+                    compilation_prompt
+                ).input_ids
+                compilation_input_ids = torch.tensor(
+                    compilation_input_ids
+                ).reshape([1, 19])
+                firstVicunaCompileInput = (compilation_input_ids,)
+                model = FirstVicuna(self.hf_model_path)
+
+                print(f"[DEBUG] generating torchscript graph")
+                ts_graph = import_with_fx(
+                    model,
+                    firstVicunaCompileInput,
+                    is_f16=self.precision == "fp16",
+                    f16_input_mask=[False, False],
+                    mlir_type="torchscript",
+                )
+                del model
+                print(f"[DEBUG] generating torch mlir")
+
+                firstVicunaCompileInput = list(firstVicunaCompileInput)
+                firstVicunaCompileInput[0] = torch_mlir.TensorPlaceholder.like(
+                    firstVicunaCompileInput[0], dynamic_axes=[1]
+                )
+                firstVicunaCompileInput = tuple(firstVicunaCompileInput)
+                module = torch_mlir.compile(
+                    ts_graph,
+                    [*firstVicunaCompileInput],
+                    torch_mlir.OutputType.LINALG_ON_TENSORS,
+                    use_tracing=False,
+                    verbose=False,
+                )
+                del ts_graph
+
+                def remove_constant_dim(line):
+                    if "19x" in line:
+                        line = re.sub("19x", "?x", line)
+                        line = re.sub(
+                            "tensor.empty\(\)", "tensor.empty(%dim)", line
+                        )
+                    if "tensor.empty" in line and "?x?" in line:
+                        line = re.sub(
+                            "tensor.empty\(%dim\)",
+                            "tensor.empty(%dim, %dim)",
+                            line,
+                        )
+                    if "arith.cmpi" in line:
+                        line = re.sub("c19", "dim", line)
+                    if " 19," in line:
+                        line = re.sub(" 19,", " %dim,", line)
+                    return line
+
+                module = str(module)
+                new_lines = []
+
+                print(f"[DEBUG] rewriting torch_mlir file")
+                for line in module.splitlines():
+                    line = remove_constant_dim(line)
+                    if "%0 = tensor.empty(%dim) : tensor<?xi64>" in line:
+                        new_lines.append(
+                            "%dim = tensor.dim %arg0, %c1 : tensor<1x?xi64>"
+                        )
+                    if (
+                        "%dim = tensor.dim %arg0, %c1 : tensor<1x?xi64>"
+                        in line
+                    ):
+                        continue
+
+                    new_lines.append(line)
+
+                module = "\n".join(new_lines)
+
+                print(f"[DEBUG] converting to bytecode")
+                del new_lines
+                module = module.encode("UTF-8")
+                module = BytesIO(module)
+                bytecode = module.read()
+                del module
+
+                print(f"[DEBUG] writing mlir to file")
+                f_ = open(self.first_vicuna_mlir_path, "wb")
+                f_.write(bytecode)
+                f_.close()
+
+        shark_module = SharkInference(
+            mlir_module=bytecode, device=self.device, mlir_dialect="tm_tensor"
+        )
+        path = shark_module.save_module(
+            self.first_vicuna_vmfb_path.parent.absolute(),
+            self.first_vicuna_vmfb_path.stem,
+            extra_args=[
+                "--iree-hal-dump-executable-sources-to=ies",
+                "--iree-vm-target-truncate-unsupported-floats",
+                "--iree-codegen-check-ir-before-llvm-conversion=false",
+                "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
+            ],
+        )
+        print("Saved first vic vmfb at vmfb at ", str(path))
+        shark_module.load_module(path)
+
+        return shark_module
+
+    def compile_second_vicuna(self):
+        vmfb = get_vmfb_from_path(
+            self.second_vicuna_vmfb_path, self.device, "tm_tensor"
+        )
+        if vmfb is not None:
+            return vmfb
+
+        # Compilation path needs some more work before it is functional
+        print(
+            f"[DEBUG] mlir path {self.second_vicuna_mlir_path} {'exists' if self.second_vicuna_mlir_path.exists() else 'does not exist'}"
+        )
+        if self.second_vicuna_mlir_path.exists():
+            with open(self.second_vicuna_mlir_path, "rb") as f:
+                bytecode = f.read()
+        else:
+            mlir_generated = False
+            if self.load_mlir_from_shark_tank:
+                if self.precision == "fp32":
+                    # download MLIR from shark_tank for fp32
+                    download_public_file(
+                        "gs://shark_tank/vicuna/unsharded/mlir/second_vicuna.mlir",
+                        self.second_vicuna_mlir_path.absolute(),
+                        single_file=True,
+                    )
+                    if self.second_vicuna_mlir_path.exists():
+                        with open(self.second_vicuna_mlir_path, "rb") as f:
+                            bytecode = f.read()
+                        mlir_generated = True
+                    else:
+                        raise ValueError(
+                            f"MLIR not found at {self.second_vicuna_mlir_path.absolute()}"
+                            " after downloading! Please check path and try again"
+                        )
+                else:
+                    print(
+                        "Only fp32 mlir added to tank, generating mlir on device."
+                    )
+
+            if not mlir_generated:
+                compilation_input_ids = torch.zeros([1, 1], dtype=torch.int64)
+                pkv = tuple(
+                    (torch.zeros([1, 32, 19, 128], dtype=torch.float32))
+                    for _ in range(64)
+                )
+                secondVicunaCompileInput = (compilation_input_ids,) + pkv
+                model = SecondVicuna(self.hf_model_path)
+                ts_graph = import_with_fx(
+                    model,
+                    secondVicunaCompileInput,
+                    is_f16=self.precision == "fp16",
+                    f16_input_mask=[False, False],
+                    mlir_type="torchscript",
+                )
+                secondVicunaCompileInput = list(secondVicunaCompileInput)
+                for i in range(len(secondVicunaCompileInput)):
+                    if i != 0:
+                        secondVicunaCompileInput[
+                            i
+                        ] = torch_mlir.TensorPlaceholder.like(
+                            secondVicunaCompileInput[i], dynamic_axes=[2]
+                        )
+                secondVicunaCompileInput = tuple(secondVicunaCompileInput)
+                module = torch_mlir.compile(
+                    ts_graph,
+                    [*secondVicunaCompileInput],
+                    torch_mlir.OutputType.LINALG_ON_TENSORS,
+                    use_tracing=False,
+                    verbose=False,
+                )
+
+                def remove_constant_dim(line):
+                    if "c19_i64" in line:
+                        line = re.sub("c19_i64", "dim_i64", line)
+                    if "19x" in line:
+                        line = re.sub("19x", "?x", line)
+                        line = re.sub(
+                            "tensor.empty\(\)", "tensor.empty(%dim)", line
+                        )
+                    if "tensor.empty" in line and "?x?" in line:
+                        line = re.sub(
+                            "tensor.empty\(%dim\)",
+                            "tensor.empty(%dim, %dim)",
+                            line,
+                        )
+                    if "arith.cmpi" in line:
+                        line = re.sub("c19", "dim", line)
+                    if " 19," in line:
+                        line = re.sub(" 19,", " %dim,", line)
+                    if "20x" in line:
+                        line = re.sub("20x", "?x", line)
+                        line = re.sub(
+                            "tensor.empty\(\)", "tensor.empty(%dimp1)", line
+                        )
+                    if " 20," in line:
+                        line = re.sub(" 20,", " %dimp1,", line)
+                    return line
+
+                module_str = str(module)
+                new_lines = []
+
+                for line in module_str.splitlines():
+                    if "%c19_i64 = arith.constant 19 : i64" in line:
+                        new_lines.append("%c2 = arith.constant 2 : index")
+                        new_lines.append(
+                            "%dim_4_int = tensor.dim %arg1, %c2 : tensor<1x32x?x128xf32>"
+                        )
+                        new_lines.append(
+                            "%dim_i64 = arith.index_cast %dim_4_int : index to i64"
+                        )
+                        continue
+                    if "%c2 = arith.constant 2 : index" in line:
+                        continue
+                    if "%c20_i64 = arith.constant 20 : i64" in line:
+                        new_lines.append("%c1_i64 = arith.constant 1 : i64")
+                        new_lines.append(
+                            "%c20_i64 = arith.addi %dim_i64, %c1_i64 : i64"
+                        )
+                        new_lines.append(
+                            "%dimp1 = arith.index_cast %c20_i64 : i64 to index"
+                        )
+                        continue
+                    line = remove_constant_dim(line)
+                    new_lines.append(line)
+
+                module_str = "\n".join(new_lines)
+                bytecode = module_str.encode("UTF-8")
+                bytecode_stream = BytesIO(bytecode)
+                bytecode = bytecode_stream.read()
+                f_ = open(self.second_vicuna_mlir_path, "wb")
+                f_.write(bytecode)
+                f_.close()
+
+        shark_module = SharkInference(
+            mlir_module=bytecode, device=self.device, mlir_dialect="tm_tensor"
+        )
+
+        path = shark_module.save_module(
+            self.second_vicuna_vmfb_path.parent.absolute(),
+            self.second_vicuna_vmfb_path.stem,
+            extra_args=[
+                "--iree-hal-dump-executable-sources-to=ies",
+                "--iree-vm-target-truncate-unsupported-floats",
+                "--iree-codegen-check-ir-before-llvm-conversion=false",
+                "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
+            ],
+        )
+        print("Saved vmfb at ", str(path))
+        shark_module.load_module(self.second_vicuna_vmfb_path)
+
+        # self.shark_module = shark_module
+
+        return shark_module
+
+    def compile(self):
+        # Cannot load both the models in the memory at once
+        # due to memory constraints, hence on demand compilation
+        # is being used until the space is enough for both models
+
+        # Testing : DO NOT Download Vmfbs if not found. Modify later
+        # download vmfbs for A100
+        if (
+            not self.first_vicuna_vmfb_path.exists()
+            and self.device == "cuda"
+            and self.precision == "fp32"
+        ):
+            download_public_file(
+                "gs://shark_tank/vicuna/unsharded/first_vicuna.vmfb",
+                self.first_vicuna_vmfb_path.absolute(),
+                single_file=True,
+            )
+        else:
+            # get first vic
+            # TODO: Remove after testing to avoid memory overload
+            # fvic_shark_model = self.compile_first_vicuna()
+            pass
+        if (
+            not self.second_vicuna_vmfb_path.exists()
+            and self.device == "cuda"
+            and self.precision == "fp32"
+        ):
+            download_public_file(
+                "gs://shark_tank/vicuna/unsharded/second_vicuna.vmfb",
+                self.second_vicuna_vmfb_path.absolute(),
+                single_file=True,
+            )
+        else:
+            # get second vic
+            # TODO: Remove after testing to avoid memory overload
+            # svic_shark_model = self.compile_second_vicuna()
+            pass
+
+        # get first vic
+        # fvic_shark_model = self.compile_first_vicuna()
+        # get second vic
+        # svic_shark_model = self.compile_second_vicuna()
+        # return tuple of shark_modules
+        # return fvic_shark_model, svic_shark_model
+        return None
+        # return tuple of shark_modules once mem is supported
+        # return fvic_shark_model, svic_shark_model
+
+    def generate(self, prompt, cli=False):
+        # TODO: refactor for cleaner integration
+        import gc
+
+        res = []
+        res_tokens = []
+        params = {
+            "prompt": prompt,
+            "is_first": True,
+            "fv": self.compile_first_vicuna(),
+        }
+
+        generated_token_op = self.generate_new_token(params=params)
+
+        token = generated_token_op["token"]
+        logits = generated_token_op["logits"]
+        pkv = generated_token_op["pkv"]
+        detok = generated_token_op["detok"]
+
+        res.append(detok)
+        res_tokens.append(token)
+        if cli:
+            print(f"Assistant: {detok}", end=" ", flush=True)
+
+        # Clear First Vic from Memory (main and cuda)
+        del params
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        sec_vic = self.compile_second_vicuna()
+        for _ in range(self.max_num_tokens - 2):
+            params = {
+                "prompt": None,
+                "is_first": False,
+                "logits": logits,
+                "pkv": pkv,
+                "sv": sec_vic,
+            }
+
+            generated_token_op = self.generate_new_token(params=params)
+
+            token = generated_token_op["token"]
+            logits = generated_token_op["logits"]
+            pkv = generated_token_op["pkv"]
+            detok = generated_token_op["detok"]
+
+            if token == 2:
+                break
+            res_tokens.append(token)
+            if detok == "<0x0A>":
+                res.append("\n")
+                if cli:
+                    print("\n", end="", flush=True)
+            else:
+                res.append(detok)
+                if cli:
+                    print(f"{detok}", end=" ", flush=True)
+        del sec_vic, pkv, logits
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        for i in range(len(res_tokens)):
+            if type(res_tokens[i]) != int:
+                res_tokens[i] = int(res_tokens[i][0])
+
+        res_str = self.tokenizer.decode(res_tokens)
+        # print(f"[DEBUG] final output : \n{res_str}")
+        return res_str
+
+    def generate_new_token(self, params, debug=False):
+        def forward_first(first_vic, prompt, cache_outputs=False):
+            input_ids = self.tokenizer(prompt).input_ids
+            input_id_len = len(input_ids)
+            input_ids = torch.tensor(input_ids)
+            input_ids = input_ids.reshape([1, input_id_len])
+            firstVicunaInput = (input_ids,)
+            assert first_vic is not None
+            output_first_vicuna = first_vic("forward", firstVicunaInput)
+            output_first_vicuna_tensor = torch.tensor(output_first_vicuna[1:])
+            logits_first_vicuna = torch.tensor(output_first_vicuna[0])
+            if cache_outputs:
+                torch.save(
+                    logits_first_vicuna, "logits_first_vicuna_tensor.pt"
+                )
+                torch.save(
+                    output_first_vicuna_tensor, "output_first_vicuna_tensor.pt"
+                )
+            token = torch.argmax(
+                torch.tensor(logits_first_vicuna)[:, -1, :], dim=1
+            )
+            return token, logits_first_vicuna, output_first_vicuna_tensor
+
+        def forward_second(sec_vic, inputs=None, load_inputs=False):
+            if inputs is not None:
+                logits = inputs[0]
+                pkv = inputs[1:]
+            elif load_inputs:
+                pkv = torch.load("output_first_vicuna_tensor.pt")
+                pkv = tuple(torch.tensor(x) for x in pkv)
+                logits = torch.load("logits_first_vicuna_tensor.pt")
+            else:
+                print(
+                    "Either inputs must be given, or load_inputs must be true"
+                )
+                return None
+            token = torch.argmax(torch.tensor(logits)[:, -1, :], dim=1)
+            token = token.to(torch.int64).reshape([1, 1])
+            secondVicunaInput = (token,) + tuple(pkv)
+
+            secondVicunaOutput = sec_vic("forward", secondVicunaInput)
+            new_pkv = secondVicunaOutput[1:]
+            new_logits = secondVicunaOutput[0]
+            new_token = torch.argmax(torch.tensor(new_logits)[:, -1, :], dim=1)
+            return new_token, new_logits, new_pkv
+
+        is_first = params["is_first"]
+
+        if is_first:
+            prompt = params["prompt"]
+            fv = params["fv"]
+            token, logits, pkv = forward_first(
+                fv,  # self.shark_model[0],
+                prompt=prompt,
+                cache_outputs=False,
+            )
+        else:
+            _logits = params["logits"]
+            _pkv = params["pkv"]
+            inputs = (_logits,) + tuple(_pkv)
+            sv = params["sv"]
+            token, logits, pkv = forward_second(
+                sv,  # self.shark_model[1],
+                inputs=inputs,
+                load_inputs=False,
+            )
+
+        detok = self.tokenizer.decode(token)
+        if debug:
+            print(
+                f"[DEBUG] is_first: {is_first} |"
+                f" token : {token} | detok : {detok}"
+            )
+        ret_dict = {
+            "token": token,
+            "logits": logits,
+            "pkv": pkv,
+            "detok": detok,
+        }
+        return ret_dict
+
+    def autocomplete(self, prompt):
+        # use First vic alone to complete a story / prompt / sentence.
+        pass

apps/language_models/src/pipelines/vicuna_sharded_pipeline.py

@@ -0,0 +1,402 @@
+from apps.language_models.src.model_wrappers.vicuna_sharded_model import (
+    FirstVicunaLayer,
+    SecondVicunaLayer,
+    CompiledFirstVicunaLayer,
+    CompiledSecondVicunaLayer,
+    ShardedVicunaModel,
+)
+from apps.language_models.src.pipelines.SharkLLMBase import SharkLLMBase
+from shark.shark_importer import import_with_fx
+from io import BytesIO
+from pathlib import Path
+from shark.shark_inference import SharkInference
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from tqdm import tqdm
+from torch_mlir import TensorPlaceholder
+
+
+import re
+import torch
+import torch_mlir
+import os
+
+
+class Vicuna(SharkLLMBase):
+    def __init__(
+        self,
+        model_name,
+        hf_model_path="TheBloke/vicuna-7B-1.1-HF",
+        max_num_tokens=512,
+        device="cuda",
+        precision="fp32",
+    ) -> None:
+        super().__init__(model_name, hf_model_path, max_num_tokens)
+        self.max_sequence_length = 256
+        self.device = device
+        self.precision = precision
+        self.tokenizer = self.get_tokenizer()
+        self.shark_model = self.compile(device=device)
+
+    def get_tokenizer(self):
+        tokenizer = AutoTokenizer.from_pretrained(
+            self.hf_model_path, use_fast=False
+        )
+        return tokenizer
+
+    def get_src_model(self):
+        kwargs = {"torch_dtype": torch.float}
+        vicuna_model = AutoModelForCausalLM.from_pretrained(
+            self.hf_model_path, **kwargs
+        )
+        return vicuna_model
+
+    def write_in_dynamic_inputs0(self, module, dynamic_input_size):
+        new_lines = []
+        for line in module.splitlines():
+            line = re.sub(f"{dynamic_input_size}x", "?x", line)
+            if "?x" in line:
+                line = re.sub("tensor.empty\(\)", "tensor.empty(%dim)", line)
+            line = re.sub(f" {dynamic_input_size},", " %dim,", line)
+            if "tensor.empty" in line and "?x?" in line:
+                line = re.sub(
+                    "tensor.empty\(%dim\)", "tensor.empty(%dim, %dim)", line
+                )
+            if "arith.cmpi" in line:
+                line = re.sub(f"c{dynamic_input_size}", "dim", line)
+            new_lines.append(line)
+        new_module = "\n".join(new_lines)
+        return new_module
+
+    def write_in_dynamic_inputs1(self, module, dynamic_input_size):
+        new_lines = []
+        for line in module.splitlines():
+            if "dim_42 =" in line:
+                continue
+            if f"%c{dynamic_input_size}_i64 =" in line:
+                new_lines.append(
+                    "%dim_42 = tensor.dim %arg1, %c3 : tensor<1x1x1x?xf32>"
+                )
+                new_lines.append(
+                    f"%dim_42_i64 = arith.index_cast %dim_42 : index to i64"
+                )
+                continue
+            line = re.sub(f"{dynamic_input_size}x", "?x", line)
+            if "?x" in line:
+                line = re.sub(
+                    "tensor.empty\(\)", "tensor.empty(%dim_42)", line
+                )
+            line = re.sub(f" {dynamic_input_size},", " %dim_42,", line)
+            if "tensor.empty" in line and "?x?" in line:
+                line = re.sub(
+                    "tensor.empty\(%dim_42\)",
+                    "tensor.empty(%dim_42, %dim_42)",
+                    line,
+                )
+            if "arith.cmpi" in line:
+                line = re.sub(f"c{dynamic_input_size}", "dim_42", line)
+            new_lines.append(line)
+        new_module = "\n".join(new_lines)
+        return new_module
+
+    def compile_vicuna_layer(
+        self,
+        vicuna_layer,
+        hidden_states,
+        attention_mask,
+        position_ids,
+        past_key_value0=None,
+        past_key_value1=None,
+    ):
+        if past_key_value0 is None and past_key_value1 is None:
+            model_inputs = (hidden_states, attention_mask, position_ids)
+        else:
+            model_inputs = (
+                hidden_states,
+                attention_mask,
+                position_ids,
+                past_key_value0,
+                past_key_value1,
+            )
+        mlir_bytecode = import_with_fx(
+            vicuna_layer,
+            model_inputs,
+            is_f16=self.precision == "fp16",
+            f16_input_mask=[False, False],
+            mlir_type="torchscript",
+        )
+        return mlir_bytecode
+
+    def compile_to_vmfb(self, inputs, layers, device="cpu", is_first=True):
+        mlirs, modules = [], []
+        for idx, layer in tqdm(enumerate(layers), desc="Getting mlirs"):
+            if is_first:
+                mlir_path = Path(f"{idx}_0.mlir")
+                vmfb_path = Path(f"{idx}_0.vmfb")
+            else:
+                mlir_path = Path(f"{idx}_1.mlir")
+                vmfb_path = Path(f"{idx}_1.vmfb")
+            if vmfb_path.exists():
+                continue
+            if mlir_path.exists():
+                # print(f"Found layer {idx} mlir")
+                f_ = open(mlir_path, "rb")
+                bytecode = f_.read()
+                f_.close()
+            else:
+                hidden_states_placeholder = TensorPlaceholder.like(
+                    inputs[0], dynamic_axes=[1]
+                )
+                attention_mask_placeholder = TensorPlaceholder.like(
+                    inputs[1], dynamic_axes=[3]
+                )
+                position_ids_placeholder = TensorPlaceholder.like(
+                    inputs[2], dynamic_axes=[1]
+                )
+                if not is_first:
+                    pkv0_placeholder = TensorPlaceholder.like(
+                        inputs[3], dynamic_axes=[2]
+                    )
+                    pkv1_placeholder = TensorPlaceholder.like(
+                        inputs[4], dynamic_axes=[2]
+                    )
+                print(f"Compiling layer {idx} mlir")
+                if is_first:
+                    ts_g = self.compile_vicuna_layer(
+                        layer, inputs[0], inputs[1], inputs[2]
+                    )
+                    module = torch_mlir.compile(
+                        ts_g,
+                        (
+                            hidden_states_placeholder,
+                            inputs[1],
+                            inputs[2],
+                        ),
+                        torch_mlir.OutputType.LINALG_ON_TENSORS,
+                        use_tracing=False,
+                        verbose=False,
+                    )
+                else:
+                    ts_g = self.compile_vicuna_layer(
+                        layer,
+                        inputs[0],
+                        inputs[1],
+                        inputs[2],
+                        inputs[3],
+                        inputs[4],
+                    )
+                    module = torch_mlir.compile(
+                        ts_g,
+                        (
+                            inputs[0],
+                            attention_mask_placeholder,
+                            inputs[2],
+                            pkv0_placeholder,
+                            pkv1_placeholder,
+                        ),
+                        torch_mlir.OutputType.LINALG_ON_TENSORS,
+                        use_tracing=False,
+                        verbose=False,
+                    )
+
+                # bytecode_stream = BytesIO()
+                # module.operation.write_bytecode(bytecode_stream)
+                # bytecode = bytecode_stream.getvalue()
+
+                if is_first:
+                    module = self.write_in_dynamic_inputs0(str(module), 137)
+                    bytecode = module.encode("UTF-8")
+                    bytecode_stream = BytesIO(bytecode)
+                    bytecode = bytecode_stream.read()
+
+                else:
+                    module = self.write_in_dynamic_inputs1(str(module), 138)
+
+                    bytecode = module.encode("UTF-8")
+                    bytecode_stream = BytesIO(bytecode)
+                    bytecode = bytecode_stream.read()
+
+                f_ = open(mlir_path, "wb")
+                f_.write(bytecode)
+                f_.close()
+            mlirs.append(bytecode)
+
+        for idx, layer in tqdm(enumerate(layers), desc="compiling modules"):
+            if is_first:
+                vmfb_path = Path(f"{idx}_0.vmfb")
+                if vmfb_path.exists():
+                    # print(f"Found layer {idx} vmfb")
+                    module = SharkInference(
+                        None,
+                        device=device,
+                        device_idx=idx % 1,
+                        mlir_dialect="tm_tensor",
+                    )
+                    module.load_module(vmfb_path)
+                else:
+                    print(f"Compiling layer {idx} vmfb")
+                    module = SharkInference(
+                        mlirs[idx],
+                        device=device,
+                        device_idx=idx % 1,
+                        mlir_dialect="tm_tensor",
+                    )
+                    module.save_module(
+                        module_name=f"{idx}_0",
+                        extra_args=[
+                            "--iree-hal-dump-executable-sources-to=ies",
+                            "--iree-vm-target-truncate-unsupported-floats",
+                            "--iree-codegen-check-ir-before-llvm-conversion=false",
+                            "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
+                        ],
+                    )
+                    module.load_module(vmfb_path)
+                modules.append(module)
+            else:
+                vmfb_path = Path(f"{idx}_1.vmfb")
+                if vmfb_path.exists():
+                    # print(f"Found layer {idx} vmfb")
+                    module = SharkInference(
+                        None,
+                        device=device,
+                        device_idx=idx % 1,
+                        mlir_dialect="tm_tensor",
+                    )
+                    module.load_module(vmfb_path)
+                else:
+                    print(f"Compiling layer {idx} vmfb")
+                    module = SharkInference(
+                        mlirs[idx],
+                        device=device,
+                        device_idx=idx % 1,
+                        mlir_dialect="tm_tensor",
+                    )
+                    module.save_module(
+                        module_name=f"{idx}_1",
+                        extra_args=[
+                            "--iree-hal-dump-executable-sources-to=ies",
+                            "--iree-vm-target-truncate-unsupported-floats",
+                            "--iree-codegen-check-ir-before-llvm-conversion=false",
+                            "--iree-vm-bytecode-module-output-format=flatbuffer-binary",
+                        ],
+                    )
+                    module.load_module(vmfb_path)
+                modules.append(module)
+
+        return mlirs, modules
+
+    def get_sharded_model(self, device="cpu"):
+        # SAMPLE_INPUT_LEN is used for creating mlir with dynamic inputs, which is currently an increadibly hacky proccess
+        # please don't change it
+        SAMPLE_INPUT_LEN = 137
+        vicuna_model = self.get_src_model()
+        placeholder_input0 = (
+            torch.zeros([1, SAMPLE_INPUT_LEN, 4096]),
+            torch.zeros([1, 1, SAMPLE_INPUT_LEN, SAMPLE_INPUT_LEN]),
+            torch.zeros([1, SAMPLE_INPUT_LEN], dtype=torch.int64),
+        )
+
+        placeholder_input1 = (
+            torch.zeros([1, 1, 4096]),
+            torch.zeros([1, 1, 1, SAMPLE_INPUT_LEN + 1]),
+            torch.zeros([1, 1], dtype=torch.int64),
+            torch.zeros([1, 32, SAMPLE_INPUT_LEN, 128]),
+            torch.zeros([1, 32, SAMPLE_INPUT_LEN, 128]),
+        )
+
+        layers0 = [
+            FirstVicunaLayer(layer) for layer in vicuna_model.model.layers
+        ]
+        _, modules0 = self.compile_to_vmfb(
+            placeholder_input0,
+            layers0,
+            is_first=True,
+            device=device,
+        )
+        shark_layers0 = [CompiledFirstVicunaLayer(m) for m in modules0]
+
+        layers1 = [
+            SecondVicunaLayer(layer) for layer in vicuna_model.model.layers
+        ]
+        _, modules1 = self.compile_to_vmfb(
+            placeholder_input1, layers1, is_first=False, device=device
+        )
+        shark_layers1 = [CompiledSecondVicunaLayer(m) for m in modules1]
+
+        sharded_model = ShardedVicunaModel(
+            vicuna_model, shark_layers0, shark_layers1
+        )
+        return sharded_model
+
+    def compile(self, device="cpu"):
+        return self.get_sharded_model(device=device)
+
+    def generate(self, prompt, cli=False):
+        # TODO: refactor for cleaner integration
+
+        tokens_generated = []
+        _past_key_values = None
+        _token = None
+        detoks_generated = []
+        for iteration in range(self.max_num_tokens):
+            params = {
+                "prompt": prompt,
+                "is_first": iteration == 0,
+                "token": _token,
+                "past_key_values": _past_key_values,
+            }
+
+            generated_token_op = self.generate_new_token(params=params)
+
+            _token = generated_token_op["token"]
+            _past_key_values = generated_token_op["past_key_values"]
+            _detok = generated_token_op["detok"]
+
+            if _token == 2:
+                break
+            detoks_generated.append(_detok)
+            tokens_generated.append(_token)
+
+        for i in range(len(tokens_generated)):
+            if type(tokens_generated[i]) != int:
+                tokens_generated[i] = int(tokens_generated[i][0])
+        result_output = self.tokenizer.decode(tokens_generated)
+        return result_output
+
+    def generate_new_token(self, params):
+        is_first = params["is_first"]
+        if is_first:
+            prompt = params["prompt"]
+            input_ids = self.tokenizer(prompt).input_ids
+            input_id_len = len(input_ids)
+            input_ids = torch.tensor(input_ids)
+            input_ids = input_ids.reshape([1, input_id_len])
+            output = self.shark_model.forward(input_ids, is_first=is_first)
+        else:
+            token = params["token"]
+            past_key_values = params["past_key_values"]
+            input_ids = [token]
+            input_id_len = len(input_ids)
+            input_ids = torch.tensor(input_ids)
+            input_ids = input_ids.reshape([1, input_id_len])
+            output = self.shark_model.forward(
+                input_ids, past_key_values=past_key_values, is_first=is_first
+            )
+
+        _logits = output["logits"]
+        _past_key_values = output["past_key_values"]
+        _token = int(torch.argmax(_logits[:, -1, :], dim=1)[0])
+        _detok = self.tokenizer.decode(_token)
+
+        ret_dict = {
+            "token": _token,
+            "detok": _detok,
+            "past_key_values": _past_key_values,
+        }
+
+        print(f" token : {_token} | detok : {_detok}")
+
+        return ret_dict
+
+    def autocomplete(self, prompt):
+        # use First vic alone to complete a story / prompt / sentence.
+        pass

apps/language_models/utils.py

@@ -3,12 +3,11 @@ from torch.fx.experimental.proxy_tensor import make_fx
 from torch._decomp import get_decompositions
 from typing import List
 from pathlib import Path
-from shark.shark_downloader import download_public_file
 
 
 # expects a Path / str as arg
 # returns None if path not found or SharkInference module
-def get_vmfb_from_path(vmfb_path, device, mlir_dialect, device_id=None):
+def get_vmfb_from_path(vmfb_path, device, mlir_dialect):
     if not isinstance(vmfb_path, Path):
         vmfb_path = Path(vmfb_path)
 
@@ -18,31 +17,9 @@ def get_vmfb_from_path(vmfb_path, device, mlir_dialect, device_id=None):
         return None
 
     print("Loading vmfb from: ", vmfb_path)
-    print("Device from get_vmfb_from_path - ", device)
     shark_module = SharkInference(
-        None, device=device, mlir_dialect=mlir_dialect, device_idx=device_id
+        None, device=device, mlir_dialect=mlir_dialect
     )
     shark_module.load_module(vmfb_path)
     print("Successfully loaded vmfb")
     return shark_module
-
-
-def get_vmfb_from_config(
-    shark_container,
-    model,
-    precision,
-    device,
-    vmfb_path,
-    padding=None,
-    device_id=None,
-):
-    vmfb_url = (
-        f"gs://shark_tank/{shark_container}/{model}_{precision}_{device}"
-    )
-    if padding:
-        vmfb_url = vmfb_url + f"_{padding}"
-    vmfb_url = vmfb_url + ".vmfb"
-    download_public_file(vmfb_url, vmfb_path.absolute(), single_file=True)
-    return get_vmfb_from_path(
-        vmfb_path, device, "tm_tensor", device_id=device_id
-    )

apps/shark_studio/api/llm.py

@@ -1,91 +0,0 @@
-from turbine_models.custom_models import stateless_llama
-from shark.iree_utils.compile_utils import get_iree_compiled_module
-from apps.shark_studio.api.utils import get_resource_path
-import iree.runtime as ireert
-import gc
-import torch
-
-llm_model_map = {
-    "llama2_7b": {
-        "initializer": stateless_llama.export_transformer_model,
-        "hf_model_name": "meta-llama/Llama-2-7b-chat-hf",
-        "stop_token": 2,
-        "max_tokens": 4096,
-    }
-}
-
-
-class LanguageModel:
-    def __init__(
-        self, model_name, hf_auth_token=None, device=None, precision="fp32"
-    ):
-        print(llm_model_map[model_name])
-        self.hf_model_name = llm_model_map[model_name]["hf_model_name"]
-        self.torch_ir, self.tokenizer = llm_model_map[model_name][
-            "initializer"
-        ](self.hf_model_name, hf_auth_token, compile_to="torch")
-        self.tempfile_name = get_resource_path("llm.torch.tempfile")
-        with open(self.tempfile_name, "w+") as f:
-            f.write(self.torch_ir)
-        del self.torch_ir
-        gc.collect()
-
-        self.device = device
-        self.precision = precision
-        self.max_tokens = llm_model_map[model_name]["max_tokens"]
-        self.iree_module_dict = None
-        self.compile()
-
-    def compile(self) -> None:
-        # this comes with keys: "vmfb", "config", and "temp_file_to_unlink".
-        self.iree_module_dict = get_iree_compiled_module(
-            self.tempfile_name, device=self.device, frontend="torch"
-        )
-        # TODO: delete the temp file
-
-    def chat(self, prompt):
-        history = []
-        for iter in range(self.max_tokens):
-            input_tensor = self.tokenizer(
-                prompt, return_tensors="pt"
-            ).input_ids
-            device_inputs = [
-                ireert.asdevicearray(
-                    self.iree_module_dict["config"], input_tensor
-                )
-            ]
-            if iter == 0:
-                token = torch.tensor(
-                    self.iree_module_dict["vmfb"]["run_initialize"](
-                        *device_inputs
-                    ).to_host()[0][0]
-                )
-            else:
-                token = torch.tensor(
-                    self.iree_module_dict["vmfb"]["run_forward"](
-                        *device_inputs
-                    ).to_host()[0][0]
-                )
-
-            history.append(token)
-            yield self.tokenizer.decode(history)
-
-            if token == llm_model_map["llama2_7b"]["stop_token"]:
-                break
-
-        for i in range(len(history)):
-            if type(history[i]) != int:
-                history[i] = int(history[i])
-        result_output = self.tokenizer.decode(history)
-        yield result_output
-
-
-if __name__ == "__main__":
-    lm = LanguageModel(
-        "llama2_7b",
-        hf_auth_token="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",
-        device="cpu-task",
-    )
-    print("model loaded")
-    for i in lm.chat("Hello, I am a robot."):
-        print(i)

apps/shark_studio/api/utils.py

@@ -1,14 +0,0 @@
-import os
-import sys
-
-
-def get_available_devices():
-    return ["cpu-task"]
-
-
-def get_resource_path(relative_path):
-    """Get absolute path to resource, works for dev and for PyInstaller"""
-    base_path = getattr(
-        sys, "_MEIPASS", os.path.dirname(os.path.abspath(__file__))
-    )
-    return os.path.join(base_path, relative_path)

apps/shark_studio/web/index.py

@@ -1,428 +0,0 @@
-from multiprocessing import Process, freeze_support
-import os
-import sys
-import logging
-from ui.chat import chat_element
-
-if sys.platform == "darwin":
-    os.environ["DYLD_LIBRARY_PATH"] = "/usr/local/lib"
-    # import before IREE to avoid MLIR library issues
-    import torch_mlir
-
-# import PIL, transformers, sentencepiece  # ensures inclusion in pysintaller exe generation
-# from apps.stable_diffusion.src import args, clear_all
-# import apps.stable_diffusion.web.utils.global_obj as global_obj
-
-
-def launch_app(address):
-    from tkinter import Tk
-    import webview
-
-    window = Tk()
-
-    # get screen width and height of display and make it more reasonably
-    # sized as we aren't making it full-screen or maximized
-    width = int(window.winfo_screenwidth() * 0.81)
-    height = int(window.winfo_screenheight() * 0.91)
-    webview.create_window(
-        "SHARK AI Studio",
-        url=address,
-        width=width,
-        height=height,
-        text_select=True,
-    )
-    webview.start(private_mode=False, storage_path=os.getcwd())
-
-
-if __name__ == "__main__":
-    # if args.debug:
-    logging.basicConfig(level=logging.DEBUG)
-    # required to do multiprocessing in a pyinstaller freeze
-    freeze_support()
-    #    if args.api or "api" in args.ui.split(","):
-    #        from apps.stable_diffusion.web.ui import (
-    #            txt2img_api,
-    #            img2img_api,
-    #            upscaler_api,
-    #            inpaint_api,
-    #            outpaint_api,
-    #            llm_chat_api,
-    #        )
-    #
-    #        from fastapi import FastAPI, APIRouter
-    #        import uvicorn
-    #
-    #        # init global sd pipeline and config
-    #        global_obj._init()
-    #
-    #        app = FastAPI()
-    #        app.add_api_route("/sdapi/v1/txt2img", txt2img_api, methods=["post"])
-    #        app.add_api_route("/sdapi/v1/img2img", img2img_api, methods=["post"])
-    #        app.add_api_route("/sdapi/v1/inpaint", inpaint_api, methods=["post"])
-    #        app.add_api_route("/sdapi/v1/outpaint", outpaint_api, methods=["post"])
-    #        app.add_api_route("/sdapi/v1/upscaler", upscaler_api, methods=["post"])
-    #
-    #        # chat APIs needed for compatibility with multiple extensions using OpenAI API
-    #        app.add_api_route(
-    #            "/v1/chat/completions", llm_chat_api, methods=["post"]
-    #        )
-    #        app.add_api_route("/v1/completions", llm_chat_api, methods=["post"])
-    #        app.add_api_route("/chat/completions", llm_chat_api, methods=["post"])
-    #        app.add_api_route("/completions", llm_chat_api, methods=["post"])
-    #        app.add_api_route(
-    #            "/v1/engines/codegen/completions", llm_chat_api, methods=["post"]
-    #        )
-    #        app.include_router(APIRouter())
-    #        uvicorn.run(app, host="0.0.0.0", port=args.server_port)
-    #        sys.exit(0)
-    #
-    # Setup to use shark_tmp for gradio's temporary image files and clear any
-    # existing temporary images there if they exist. Then we can import gradio.
-    # It has to be in this order or gradio ignores what we've set up.
-    # from apps.stable_diffusion.web.utils.gradio_configs import (
-    #    config_gradio_tmp_imgs_folder,
-    # )
-
-    # config_gradio_tmp_imgs_folder()
-    import gradio as gr
-
-    # Create custom models folders if they don't exist
-    # from apps.stable_diffusion.web.ui.utils import create_custom_models_folders
-
-    # create_custom_models_folders()
-
-    def resource_path(relative_path):
-        """Get absolute path to resource, works for dev and for PyInstaller"""
-        base_path = getattr(
-            sys, "_MEIPASS", os.path.dirname(os.path.abspath(__file__))
-        )
-        return os.path.join(base_path, relative_path)
-
-    dark_theme = resource_path("ui/css/sd_dark_theme.css")
-
-    # from apps.stable_diffusion.web.ui import (
-    # txt2img_web,
-    # txt2img_custom_model,
-    # txt2img_gallery,
-    # txt2img_png_info_img,
-    # txt2img_status,
-    # txt2img_sendto_img2img,
-    # txt2img_sendto_inpaint,
-    # txt2img_sendto_outpaint,
-    # txt2img_sendto_upscaler,
-    ## h2ogpt_upload,
-    ## h2ogpt_web,
-    # img2img_web,
-    # img2img_custom_model,
-    # img2img_gallery,
-    # img2img_init_image,
-    # img2img_status,
-    # img2img_sendto_inpaint,
-    # img2img_sendto_outpaint,
-    # img2img_sendto_upscaler,
-    # inpaint_web,
-    # inpaint_custom_model,
-    # inpaint_gallery,
-    # inpaint_init_image,
-    # inpaint_status,
-    # inpaint_sendto_img2img,
-    # inpaint_sendto_outpaint,
-    # inpaint_sendto_upscaler,
-    # outpaint_web,
-    # outpaint_custom_model,
-    # outpaint_gallery,
-    # outpaint_init_image,
-    # outpaint_status,
-    # outpaint_sendto_img2img,
-    # outpaint_sendto_inpaint,
-    # outpaint_sendto_upscaler,
-    # upscaler_web,
-    # upscaler_custom_model,
-    # upscaler_gallery,
-    # upscaler_init_image,
-    # upscaler_status,
-    # upscaler_sendto_img2img,
-    # upscaler_sendto_inpaint,
-    # upscaler_sendto_outpaint,
-    ##  lora_train_web,
-    ##  model_web,
-    ##  model_config_web,
-    # hf_models,
-    # modelmanager_sendto_txt2img,
-    # modelmanager_sendto_img2img,
-    # modelmanager_sendto_inpaint,
-    # modelmanager_sendto_outpaint,
-    # modelmanager_sendto_upscaler,
-    # stablelm_chat,
-    # minigpt4_web,
-    # outputgallery_web,
-    # outputgallery_tab_select,
-    # outputgallery_watch,
-    # outputgallery_filename,
-    # outputgallery_sendto_txt2img,
-    # outputgallery_sendto_img2img,
-    # outputgallery_sendto_inpaint,
-    # outputgallery_sendto_outpaint,
-    # outputgallery_sendto_upscaler,
-    # )
-
-    # init global sd pipeline and config
-    # global_obj._init()
-
-    def register_button_click(button, selectedid, inputs, outputs):
-        button.click(
-            lambda x: (
-                x[0]["name"] if len(x) != 0 else None,
-                gr.Tabs.update(selected=selectedid),
-            ),
-            inputs,
-            outputs,
-        )
-
-    def register_modelmanager_button(button, selectedid, inputs, outputs):
-        button.click(
-            lambda x: (
-                "None",
-                x,
-                gr.Tabs.update(selected=selectedid),
-            ),
-            inputs,
-            outputs,
-        )
-
-    def register_outputgallery_button(button, selectedid, inputs, outputs):
-        button.click(
-            lambda x: (
-                x,
-                gr.Tabs.update(selected=selectedid),
-            ),
-            inputs,
-            outputs,
-        )
-
-    with gr.Blocks(
-        css=dark_theme, analytics_enabled=False, title="Stable Diffusion"
-    ) as sd_web:
-        with gr.Tabs() as tabs:
-            # NOTE: If adding, removing, or re-ordering tabs, make sure that they
-            # have a unique id that doesn't clash with any of the other tabs,
-            # and that the order in the code here is the order they should
-            # appear in the ui, as the id value doesn't determine the order.
-
-            # Where possible, avoid changing the id of any tab that is the
-            # destination of one of the 'send to' buttons. If you do have to change
-            # that id, make sure you update the relevant register_button_click calls
-            # further down with the new id.
-            # with gr.TabItem(label="Text-to-Image", id=0):
-            #    txt2img_web.render()
-            # with gr.TabItem(label="Image-to-Image", id=1):
-            #    img2img_web.render()
-            # with gr.TabItem(label="Inpainting", id=2):
-            #    inpaint_web.render()
-            # with gr.TabItem(label="Outpainting", id=3):
-            #    outpaint_web.render()
-            # with gr.TabItem(label="Upscaler", id=4):
-            #    upscaler_web.render()
-            # if args.output_gallery:
-            #    with gr.TabItem(label="Output Gallery", id=5) as og_tab:
-            #        outputgallery_web.render()
-
-            #    # extra output gallery configuration
-            #    outputgallery_tab_select(og_tab.select)
-            #    outputgallery_watch(
-            #        [
-            #            txt2img_status,
-            #            img2img_status,
-            #            inpaint_status,
-            #            outpaint_status,
-            #            upscaler_status,
-            #        ]
-            #    )
-            ##  with gr.TabItem(label="Model Manager", id=6):
-            ##      model_web.render()
-            ##  with gr.TabItem(label="LoRA Training (Experimental)", id=7):
-            ##      lora_train_web.render()
-            with gr.TabItem(label="Chat Bot", id=0):
-                chat_element.render()
-            ##  with gr.TabItem(
-            ##      label="Generate Sharding Config (Experimental)", id=9
-            ##  ):
-            ##      model_config_web.render()
-            # with gr.TabItem(label="MultiModal (Experimental)", id=10):
-            #    minigpt4_web.render()
-            # with gr.TabItem(label="DocuChat Upload", id=11):
-            #     h2ogpt_upload.render()
-            # with gr.TabItem(label="DocuChat(Experimental)", id=12):
-            #     h2ogpt_web.render()
-
-        # send to buttons
-        # register_button_click(
-        #    txt2img_sendto_img2img,
-        #    1,
-        #    [txt2img_gallery],
-        #    [img2img_init_image, tabs],
-        # )
-        # register_button_click(
-        #    txt2img_sendto_inpaint,
-        #    2,
-        #    [txt2img_gallery],
-        #    [inpaint_init_image, tabs],
-        # )
-        # register_button_click(
-        #    txt2img_sendto_outpaint,
-        #    3,
-        #    [txt2img_gallery],
-        #    [outpaint_init_image, tabs],
-        # )
-        # register_button_click(
-        #    txt2img_sendto_upscaler,
-        #    4,
-        #    [txt2img_gallery],
-        #    [upscaler_init_image, tabs],
-        # )
-        # register_button_click(
-        #    img2img_sendto_inpaint,
-        #    2,
-        #    [img2img_gallery],
-        #    [inpaint_init_image, tabs],
-        # )
-        # register_button_click(
-        #    img2img_sendto_outpaint,
-        #    3,
-        #    [img2img_gallery],
-        #    [outpaint_init_image, tabs],
-        # )
-        # register_button_click(
-        #    img2img_sendto_upscaler,
-        #    4,
-        #    [img2img_gallery],
-        #    [upscaler_init_image, tabs],
-        # )
-        # register_button_click(
-        #    inpaint_sendto_img2img,
-        #    1,
-        #    [inpaint_gallery],
-        #    [img2img_init_image, tabs],
-        # )
-        # register_button_click(
-        #    inpaint_sendto_outpaint,
-        #    3,
-        #    [inpaint_gallery],
-        #    [outpaint_init_image, tabs],
-        # )
-        # register_button_click(
-        #    inpaint_sendto_upscaler,
-        #    4,
-        #    [inpaint_gallery],
-        #    [upscaler_init_image, tabs],
-        # )
-        # register_button_click(
-        #    outpaint_sendto_img2img,
-        #    1,
-        #    [outpaint_gallery],
-        #    [img2img_init_image, tabs],
-        # )
-        # register_button_click(
-        #    outpaint_sendto_inpaint,
-        #    2,
-        #    [outpaint_gallery],
-        #    [inpaint_init_image, tabs],
-        # )
-        # register_button_click(
-        #    outpaint_sendto_upscaler,
-        #    4,
-        #    [outpaint_gallery],
-        #    [upscaler_init_image, tabs],
-        # )
-        # register_button_click(
-        #    upscaler_sendto_img2img,
-        #    1,
-        #    [upscaler_gallery],
-        #    [img2img_init_image, tabs],
-        # )
-        # register_button_click(
-        #    upscaler_sendto_inpaint,
-        #    2,
-        #    [upscaler_gallery],
-        #    [inpaint_init_image, tabs],
-        # )
-        # register_button_click(
-        #    upscaler_sendto_outpaint,
-        #    3,
-        #    [upscaler_gallery],
-        #    [outpaint_init_image, tabs],
-        # )
-        # if args.output_gallery:
-        #    register_outputgallery_button(
-        #        outputgallery_sendto_txt2img,
-        #        0,
-        #        [outputgallery_filename],
-        #        [txt2img_png_info_img, tabs],
-        #    )
-        #    register_outputgallery_button(
-        #        outputgallery_sendto_img2img,
-        #        1,
-        #        [outputgallery_filename],
-        #        [img2img_init_image, tabs],
-        #    )
-        #    register_outputgallery_button(
-        #        outputgallery_sendto_inpaint,
-        #        2,
-        #        [outputgallery_filename],
-        #        [inpaint_init_image, tabs],
-        #    )
-        #    register_outputgallery_button(
-        #        outputgallery_sendto_outpaint,
-        #        3,
-        #        [outputgallery_filename],
-        #        [outpaint_init_image, tabs],
-        #    )
-        #    register_outputgallery_button(
-        #        outputgallery_sendto_upscaler,
-        #        4,
-        #        [outputgallery_filename],
-        #        [upscaler_init_image, tabs],
-        #    )
-        # register_modelmanager_button(
-        #    modelmanager_sendto_txt2img,
-        #    0,
-        #    [hf_models],
-        #    [txt2img_custom_model, tabs],
-        # )
-        # register_modelmanager_button(
-        #    modelmanager_sendto_img2img,
-        #    1,
-        #    [hf_models],
-        #    [img2img_custom_model, tabs],
-        # )
-        # register_modelmanager_button(
-        #    modelmanager_sendto_inpaint,
-        #    2,
-        #    [hf_models],
-        #    [inpaint_custom_model, tabs],
-        # )
-        # register_modelmanager_button(
-        #    modelmanager_sendto_outpaint,
-        #    3,
-        #    [hf_models],
-        #    [outpaint_custom_model, tabs],
-        # )
-        # register_modelmanager_button(
-        #    modelmanager_sendto_upscaler,
-        #    4,
-        #    [hf_models],
-        #    [upscaler_custom_model, tabs],
-        # )
-
-    sd_web.queue()
-    # if args.ui == "app":
-    #    t = Process(
-    #        target=launch_app, args=[f"http://localhost:{args.server_port}"]
-    #    )
-    #    t.start()
-    sd_web.launch(
-        share=True,
-        inbrowser=True,
-        server_name="0.0.0.0",
-        server_port=11911,  # args.server_port,
-    )

apps/shark_studio/web/ui/chat.py

@@ -1,517 +0,0 @@
-import gradio as gr
-import os
-from pathlib import Path
-from datetime import datetime as dt
-import json
-import sys
-from apps.shark_studio.api.utils import (
-    get_available_devices,
-)
-from apps.shark_studio.api.llm import (
-    llm_model_map,
-    LanguageModel,
-)
-
-
-def user(message, history):
-    # Append the user's message to the conversation history
-    return "", history + [[message, ""]]
-
-
-language_model = None
-
-
-# NOTE: Each `model_name` should have its own start message
-start_message = {
-    "llama2_7b": (
-        "You are a helpful, respectful and honest assistant. Always answer "
-        "as helpfully as possible, while being safe. Your answers should not "
-        "include any harmful, unethical, racist, sexist, toxic, dangerous, or "
-        "illegal content. Please ensure that your responses are socially "
-        "unbiased and positive in nature. If a question does not make any "
-        "sense, or is not factually coherent, explain why instead of "
-        "answering something not correct. If you don't know the answer "
-        "to a question, please don't share false information."
-    ),
-    "llama2_13b": (
-        "You are a helpful, respectful and honest assistant. Always answer "
-        "as helpfully as possible, while being safe. Your answers should not "
-        "include any harmful, unethical, racist, sexist, toxic, dangerous, or "
-        "illegal content. Please ensure that your responses are socially "
-        "unbiased and positive in nature. If a question does not make any "
-        "sense, or is not factually coherent, explain why instead of "
-        "answering something not correct. If you don't know the answer "
-        "to a question, please don't share false information."
-    ),
-    "llama2_70b": (
-        "You are a helpful, respectful and honest assistant. Always answer "
-        "as helpfully as possible, while being safe. Your answers should not "
-        "include any harmful, unethical, racist, sexist, toxic, dangerous, or "
-        "illegal content. Please ensure that your responses are socially "
-        "unbiased and positive in nature. If a question does not make any "
-        "sense, or is not factually coherent, explain why instead of "
-        "answering something not correct. If you don't know the answer "
-        "to a question, please don't share false information."
-    ),
-    "vicuna": (
-        "A chat between a curious user and an artificial intelligence "
-        "assistant. The assistant gives helpful, detailed, and "
-        "polite answers to the user's questions.\n"
-    ),
-}
-
-
-def create_prompt(model_name, history, prompt_prefix):
-    return ""
-    system_message = ""
-    if prompt_prefix:
-        system_message = start_message[model_name]
-
-    if "llama2" in model_name:
-        B_INST, E_INST = "[INST]", "[/INST]"
-        B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
-        conversation = "".join(
-            [f"{B_INST} {item[0]} {E_INST} {item[1]} " for item in history[1:]]
-        )
-        if prompt_prefix:
-            msg = f"{B_INST} {B_SYS}{system_message}{E_SYS}{history[0][0]} {E_INST} {history[0][1]} {conversation}"
-        else:
-            msg = f"{B_INST} {history[0][0]} {E_INST} {history[0][1]} {conversation}"
-    elif model_name in ["vicuna"]:
-        conversation = "".join(
-            [
-                "".join(["<|USER|>" + item[0], "<|ASSISTANT|>" + item[1]])
-                for item in history
-            ]
-        )
-        msg = system_message + conversation
-        msg = msg.strip()
-    else:
-        conversation = "".join(
-            ["".join([item[0], item[1]]) for item in history]
-        )
-        msg = system_message + conversation
-        msg = msg.strip()
-    return msg
-
-
-def get_default_config():
-    return False
-    import torch
-    from transformers import AutoTokenizer
-
-    hf_model_path = "TheBloke/vicuna-7B-1.1-HF"
-    tokenizer = AutoTokenizer.from_pretrained(hf_model_path, use_fast=False)
-    compilation_prompt = "".join(["0" for _ in range(17)])
-    compilation_input_ids = tokenizer(
-        compilation_prompt,
-        return_tensors="pt",
-    ).input_ids
-    compilation_input_ids = torch.tensor(compilation_input_ids).reshape(
-        [1, 19]
-    )
-    firstVicunaCompileInput = (compilation_input_ids,)
-    from apps.language_models.src.model_wrappers.vicuna_model import (
-        CombinedModel,
-    )
-    from shark.shark_generate_model_config import GenerateConfigFile
-
-    model = CombinedModel()
-    c = GenerateConfigFile(model, 1, ["gpu_id"], firstVicunaCompileInput)
-    c.split_into_layers()
-
-
-# model_vmfb_key = ""
-
-
-def chat_fn(
-    prompt_prefix,
-    history,
-    model,
-    device,
-    precision,
-    download_vmfb,
-    config_file,
-    cli=False,
-    progress=gr.Progress(),
-):
-    global language_model
-    if language_model is None:
-        language_model = LanguageModel(
-            model, device=device, precision=precision
-        )
-
-    language_model.chat(prompt_prefix)
-    return "", ""
-    global past_key_values
-    global model_vmfb_key
-
-    device_id = None
-    model_name, model_path = list(map(str.strip, model.split("=>")))
-    if "cuda" in device:
-        device = "cuda"
-    elif "sync" in device:
-        device = "cpu-sync"
-    elif "task" in device:
-        device = "cpu-task"
-    elif "vulkan" in device:
-        device_id = int(device.split("://")[1])
-        device = "vulkan"
-    elif "rocm" in device:
-        device = "rocm"
-    else:
-        print("unrecognized device")
-
-    from apps.language_models.scripts.vicuna import ShardedVicuna
-    from apps.language_models.scripts.vicuna import UnshardedVicuna
-    from apps.stable_diffusion.src import args
-
-    new_model_vmfb_key = f"{model_name}#{model_path}#{device}#{device_id}#{precision}#{download_vmfb}"
-    if vicuna_model is None or new_model_vmfb_key != model_vmfb_key:
-        model_vmfb_key = new_model_vmfb_key
-        max_toks = 128 if model_name == "codegen" else 512
-
-        # get iree flags that need to be overridden, from commandline args
-        _extra_args = []
-        # vulkan target triple
-        vulkan_target_triple = args.iree_vulkan_target_triple
-        from shark.iree_utils.vulkan_utils import (
-            get_all_vulkan_devices,
-            get_vulkan_target_triple,
-        )
-
-        if device == "vulkan":
-            vulkaninfo_list = get_all_vulkan_devices()
-            if vulkan_target_triple == "":
-                # We already have the device_id extracted via WebUI, so we directly use
-                # that to find the target triple.
-                vulkan_target_triple = get_vulkan_target_triple(
-                    vulkaninfo_list[device_id]
-                )
-            _extra_args.append(
-                f"-iree-vulkan-target-triple={vulkan_target_triple}"
-            )
-            if "rdna" in vulkan_target_triple:
-                flags_to_add = [
-                    "--iree-spirv-index-bits=64",
-                ]
-                _extra_args = _extra_args + flags_to_add
-
-            if device_id is None:
-                id = 0
-                for device in vulkaninfo_list:
-                    target_triple = get_vulkan_target_triple(
-                        vulkaninfo_list[id]
-                    )
-                    if target_triple == vulkan_target_triple:
-                        device_id = id
-                        break
-                    id += 1
-
-                assert (
-                    device_id
-                ), f"no vulkan hardware for target-triple '{vulkan_target_triple}' exists"
-            print(f"Will use vulkan target triple : {vulkan_target_triple}")
-
-        elif "rocm" in device:
-            # add iree rocm flags
-            _extra_args.append(
-                f"--iree-rocm-target-chip={args.iree_rocm_target_chip}"
-            )
-            print(f"extra args = {_extra_args}")
-
-        if model_name == "vicuna4":
-            vicuna_model = ShardedVicuna(
-                model_name,
-                hf_model_path=model_path,
-                device=device,
-                precision=precision,
-                max_num_tokens=max_toks,
-                compressed=True,
-                extra_args_cmd=_extra_args,
-            )
-        else:
-            #  if config_file is None:
-            vicuna_model = UnshardedVicuna(
-                model_name,
-                hf_model_path=model_path,
-                hf_auth_token=args.hf_auth_token,
-                device=device,
-                vulkan_target_triple=vulkan_target_triple,
-                precision=precision,
-                max_num_tokens=max_toks,
-                download_vmfb=download_vmfb,
-                load_mlir_from_shark_tank=True,
-                extra_args_cmd=_extra_args,
-                device_id=device_id,
-            )
-
-    if vicuna_model is None:
-        sys.exit("Unable to instantiate the model object, exiting.")
-
-    prompt = create_prompt(model_name, history, prompt_prefix)
-
-    partial_text = ""
-    token_count = 0
-    total_time_ms = 0.001  # In order to avoid divide by zero error
-    prefill_time = 0
-    is_first = True
-    for text, msg, exec_time in progress.tqdm(
-        vicuna_model.generate(prompt, cli=cli),
-        desc="generating response",
-    ):
-        if msg is None:
-            if is_first:
-                prefill_time = exec_time
-                is_first = False
-            else:
-                total_time_ms += exec_time
-                token_count += 1
-            partial_text += text + " "
-            history[-1][1] = partial_text
-            yield history, f"Prefill: {prefill_time:.2f}"
-        elif "formatted" in msg:
-            history[-1][1] = text
-            tokens_per_sec = (token_count / total_time_ms) * 1000
-            yield history, f"Prefill: {prefill_time:.2f} seconds\n Decode: {tokens_per_sec:.2f} tokens/sec"
-        else:
-            sys.exit(
-                "unexpected message from the vicuna generate call, exiting."
-            )
-
-    return history, ""
-
-
-def llm_chat_api(InputData: dict):
-    return None
-    print(f"Input keys : {InputData.keys()}")
-    # print(f"model : {InputData['model']}")
-    is_chat_completion_api = (
-        "messages" in InputData.keys()
-    )  # else it is the legacy `completion` api
-    # For Debugging input data from API
-    # if is_chat_completion_api:
-    #     print(f"message -> role : {InputData['messages'][0]['role']}")
-    #     print(f"message -> content : {InputData['messages'][0]['content']}")
-    # else:
-    #     print(f"prompt : {InputData['prompt']}")
-    # print(f"max_tokens : {InputData['max_tokens']}") # Default to 128 for now
-    global vicuna_model
-    model_name = (
-        InputData["model"] if "model" in InputData.keys() else "codegen"
-    )
-    model_path = llm_model_map[model_name]
-    device = "cpu-task"
-    precision = "fp16"
-    max_toks = (
-        None
-        if "max_tokens" not in InputData.keys()
-        else InputData["max_tokens"]
-    )
-    if max_toks is None:
-        max_toks = 128 if model_name == "codegen" else 512
-
-    # make it working for codegen first
-    from apps.language_models.scripts.vicuna import (
-        UnshardedVicuna,
-    )
-
-    device_id = None
-    if vicuna_model == 0:
-        if "cuda" in device:
-            device = "cuda"
-        elif "sync" in device:
-            device = "cpu-sync"
-        elif "task" in device:
-            device = "cpu-task"
-        elif "vulkan" in device:
-            device_id = int(device.split("://")[1])
-            device = "vulkan"
-        else:
-            print("unrecognized device")
-
-        vicuna_model = UnshardedVicuna(
-            model_name,
-            hf_model_path=model_path,
-            device=device,
-            precision=precision,
-            max_num_tokens=max_toks,
-            download_vmfb=True,
-            load_mlir_from_shark_tank=True,
-            device_id=device_id,
-        )
-
-    # TODO: add role dict for different models
-    if is_chat_completion_api:
-        # TODO: add funtionality for multiple messages
-        prompt = create_prompt(
-            model_name, [(InputData["messages"][0]["content"], "")]
-        )
-    else:
-        prompt = InputData["prompt"]
-    print("prompt = ", prompt)
-
-    res = vicuna_model.generate(prompt)
-    res_op = None
-    for op in res:
-        res_op = op
-
-    if is_chat_completion_api:
-        choices = [
-            {
-                "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": res_op,  # since we are yeilding the result
-                },
-                "finish_reason": "stop",  # or length
-            }
-        ]
-    else:
-        choices = [
-            {
-                "text": res_op,
-                "index": 0,
-                "logprobs": None,
-                "finish_reason": "stop",  # or length
-            }
-        ]
-    end_time = dt.now().strftime("%Y%m%d%H%M%S%f")
-    return {
-        "id": end_time,
-        "object": "chat.completion"
-        if is_chat_completion_api
-        else "text_completion",
-        "created": int(end_time),
-        "choices": choices,
-    }
-
-
-def view_json_file(file_obj):
-    content = ""
-    with open(file_obj.name, "r") as fopen:
-        content = fopen.read()
-    return content
-
-
-with gr.Blocks(title="Chat") as chat_element:
-    with gr.Row():
-        model_choices = list(llm_model_map.keys())
-        model = gr.Dropdown(
-            label="Select Model",
-            value=model_choices[0],
-            choices=model_choices,
-            allow_custom_value=True,
-        )
-        supported_devices = get_available_devices()
-        enabled = True
-        if len(supported_devices) == 0:
-            supported_devices = ["cpu-task"]
-        supported_devices = [x for x in supported_devices if "sync" not in x]
-        device = gr.Dropdown(
-            label="Device",
-            value=supported_devices[0],
-            choices=supported_devices,
-            interactive=enabled,
-            allow_custom_value=True,
-        )
-        precision = gr.Radio(
-            label="Precision",
-            value="int4",
-            choices=[
-                # "int4",
-                # "int8",
-                # "fp16",
-                "fp32",
-            ],
-            visible=False,
-        )
-        tokens_time = gr.Textbox(label="Tokens generated per second")
-        with gr.Column():
-            download_vmfb = gr.Checkbox(
-                label="Download vmfb from Shark tank if available",
-                value=True,
-                interactive=True,
-            )
-            prompt_prefix = gr.Checkbox(
-                label="Add System Prompt",
-                value=False,
-                interactive=True,
-            )
-
-    chatbot = gr.Chatbot(height=500)
-    with gr.Row():
-        with gr.Column():
-            msg = gr.Textbox(
-                label="Chat Message Box",
-                placeholder="Chat Message Box",
-                show_label=False,
-                interactive=enabled,
-                container=False,
-            )
-        with gr.Column():
-            with gr.Row():
-                submit = gr.Button("Submit", interactive=enabled)
-                stop = gr.Button("Stop", interactive=enabled)
-                clear = gr.Button("Clear", interactive=enabled)
-
-    with gr.Row(visible=False):
-        with gr.Group():
-            config_file = gr.File(
-                label="Upload sharding configuration", visible=False
-            )
-            json_view_button = gr.Button(label="View as JSON", visible=False)
-        json_view = gr.JSON(interactive=True, visible=False)
-        json_view_button.click(
-            fn=view_json_file, inputs=[config_file], outputs=[json_view]
-        )
-    submit_event = msg.submit(
-        fn=user,
-        inputs=[msg, chatbot],
-        outputs=[msg, chatbot],
-        show_progress=False,
-        queue=False,
-    ).then(
-        fn=chat_fn,
-        inputs=[
-            prompt_prefix,
-            chatbot,
-            model,
-            device,
-            precision,
-            download_vmfb,
-            config_file,
-        ],
-        outputs=[chatbot, tokens_time],
-        show_progress=False,
-        queue=True,
-    )
-    submit_click_event = submit.click(
-        fn=user,
-        inputs=[msg, chatbot],
-        outputs=[msg, chatbot],
-        show_progress=False,
-        queue=False,
-    ).then(
-        fn=chat_fn,
-        inputs=[
-            prompt_prefix,
-            chatbot,
-            model,
-            device,
-            precision,
-            download_vmfb,
-            config_file,
-        ],
-        outputs=[chatbot, tokens_time],
-        show_progress=False,
-        queue=True,
-    )
-    stop.click(
-        fn=None,
-        inputs=None,
-        outputs=None,
-        cancels=[submit_event, submit_click_event],
-        queue=False,
-    )
-    clear.click(lambda: None, None, [chatbot], queue=False)

apps/stable_diffusion/profiling_with_iree.md

@@ -7,16 +7,16 @@ Compile Commands FP32/FP16:
 
 ```shell
 Vulkan AMD: 
-iree-compile --iree-input-type=none --iree-hal-target-backends=vulkan --iree-vulkan-target-triple=rdna2-unknown-linux /path/to/input/mlir -o /path/to/output/vmfb
+iree-compile --iree-input-type=none --iree-hal-target-backends=vulkan --iree-vulkan-target-triple=rdna2-unknown-linux --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 /path/to/input/mlir -o /path/to/output/vmfb
 
 #  add --mlir-print-debuginfo --mlir-print-op-on-diagnostic=true for debug
 #  use –iree-input-type=auto or "mhlo_legacy" or "stablehlo" for TF models
 
 CUDA NVIDIA:
-iree-compile --iree-input-type=none --iree-hal-target-backends=cuda /path/to/input/mlir -o /path/to/output/vmfb
+iree-compile --iree-input-type=none --iree-hal-target-backends=cuda --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 /path/to/input/mlir -o /path/to/output/vmfb
 
 CPU:
-iree-compile --iree-input-type=none --iree-hal-target-backends=llvm-cpu /path/to/input/mlir -o /path/to/output/vmfb
+iree-compile --iree-input-type=none --iree-hal-target-backends=llvm-cpu  --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 /path/to/input/mlir -o /path/to/output/vmfb
 ```
 
 

apps/stable_diffusion/scripts/img2img.py

@@ -103,9 +103,7 @@ def main():
         dtype,
         args.use_base_vae,
         cpu_scheduling,
-        args.max_embeddings_multiples,
         use_stencil=use_stencil,
-        control_mode=args.control_mode,
     )
     total_time = time.time() - start_time
     text_output = f"prompt={args.prompts}"

apps/stable_diffusion/scripts/inpaint.py

@@ -58,8 +58,11 @@ def main():
         ondemand=args.ondemand,
     )
 
-    seeds = utils.batch_seeds(seed, args.batch_count, args.repeatable_seeds)
     for current_batch in range(args.batch_count):
+        if current_batch > 0:
+            seed = -1
+        seed = utils.sanitize_seed(seed)
+
         start_time = time.time()
         generated_imgs = inpaint_obj.generate_images(
             args.prompts,
@@ -73,12 +76,11 @@ def main():
             args.inpaint_full_res_padding,
             args.steps,
             args.guidance_scale,
-            seeds[current_batch],
+            seed,
             args.max_length,
             dtype,
             args.use_base_vae,
             cpu_scheduling,
-            args.max_embeddings_multiples,
         )
         total_time = time.time() - start_time
         text_output = f"prompt={args.prompts}"
@@ -87,10 +89,7 @@ def main():
             f"\nmodel_id={args.hf_model_id}, ckpt_loc={args.ckpt_loc}"
         )
         text_output += f"\nscheduler={args.scheduler}, device={args.device}"
-        text_output += (
-            f"\nsteps={args.steps}, guidance_scale={args.guidance_scale},"
-        )
-        text_output += f"seed={seed}, size={args.height}x{args.width}"
+        text_output += f"\nsteps={args.steps}, guidance_scale={args.guidance_scale}, seed={seed}, size={args.height}x{args.width}"
         text_output += (
             f", batch size={args.batch_size}, max_length={args.max_length}"
         )

apps/stable_diffusion/scripts/outpaint.py

@@ -51,8 +51,11 @@ def main():
         ondemand=args.ondemand,
     )
 
-    seeds = utils.batch_seeds(seed, args.batch_count, args.repeatable_seeds)
     for current_batch in range(args.batch_count):
+        if current_batch > 0:
+            seed = -1
+        seed = utils.sanitize_seed(seed)
+
         start_time = time.time()
         generated_imgs = outpaint_obj.generate_images(
             args.prompts,
@@ -71,12 +74,11 @@ def main():
             args.width,
             args.steps,
             args.guidance_scale,
-            seeds[current_batch],
+            seed,
             args.max_length,
             dtype,
             args.use_base_vae,
             cpu_scheduling,
-            args.max_embeddings_multiples,
         )
         total_time = time.time() - start_time
         text_output = f"prompt={args.prompts}"
@@ -85,10 +87,7 @@ def main():
             f"\nmodel_id={args.hf_model_id}, ckpt_loc={args.ckpt_loc}"
         )
         text_output += f"\nscheduler={args.scheduler}, device={args.device}"
-        text_output += (
-            f"\nsteps={args.steps}, guidance_scale={args.guidance_scale},"
-        )
-        text_output += f"seed={seed}, size={args.height}x{args.width}"
+        text_output += f"\nsteps={args.steps}, guidance_scale={args.guidance_scale}, seed={seed}, size={args.height}x{args.width}"
         text_output += (
             f", batch size={args.batch_size}, max_length={args.max_length}"
         )

apps/stable_diffusion/scripts/train_lora_word.py

@@ -34,7 +34,7 @@ from PIL import Image
 from tqdm.auto import tqdm
 from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
 from diffusers.loaders import AttnProcsLayers
-from diffusers.models.attention_processor import LoRAXFormersAttnProcessor
+from diffusers.models.cross_attention import LoRACrossAttnProcessor
 
 import torch_mlir
 from torch_mlir.dynamo import make_simple_dynamo_backend
@@ -223,8 +223,7 @@ def lora_train(
         if not hf_model_id:
             return (
                 None,
-                "Please provide either custom model or huggingface model ID, both must not be "
-                "empty.",
+                "Please provide either custom model or huggingface model ID, both must not be empty",
             )
         args.hf_model_id = hf_model_id
     elif ".ckpt" in custom_model or ".safetensors" in custom_model:
@@ -287,7 +286,7 @@ def lora_train(
                 block_id = int(name[len("down_blocks.")])
                 hidden_size = unet.config.block_out_channels[block_id]
 
-            lora_attn_procs[name] = LoRAXFormersAttnProcessor(
+            lora_attn_procs[name] = LoRACrossAttnProcessor(
                 hidden_size=hidden_size,
                 cross_attention_dim=cross_attention_dim,
             )

apps/stable_diffusion/scripts/tuner.py

@@ -17,10 +17,6 @@ from apps.stable_diffusion.src.models import SharkifyStableDiffusionModel
 
 
 def load_mlir_module():
-    if "upscaler" in args.hf_model_id:
-        is_upscaler = True
-    else:
-        is_upscaler = False
     sd_model = SharkifyStableDiffusionModel(
         args.hf_model_id,
         args.ckpt_loc,
@@ -31,7 +27,6 @@ def load_mlir_module():
         height=args.height,
         width=args.width,
         use_base_vae=args.use_base_vae,
-        is_upscaler=is_upscaler,
         use_tuned=False,
         low_cpu_mem_usage=args.low_cpu_mem_usage,
         return_mlir=True,

apps/stable_diffusion/scripts/txt2img.py

@@ -42,8 +42,11 @@ def main():
         ondemand=args.ondemand,
     )
 
-    seeds = utils.batch_seeds(seed, args.batch_count, args.repeatable_seeds)
     for current_batch in range(args.batch_count):
+        if current_batch > 0:
+            seed = -1
+        seed = utils.sanitize_seed(seed)
+
         start_time = time.time()
         generated_imgs = txt2img_obj.generate_images(
             args.prompts,
@@ -53,7 +56,7 @@ def main():
             args.width,
             args.steps,
             args.guidance_scale,
-            seeds[current_batch],
+            seed,
             args.max_length,
             dtype,
             args.use_base_vae,
@@ -67,12 +70,7 @@ def main():
             f"\nmodel_id={args.hf_model_id}, ckpt_loc={args.ckpt_loc}"
         )
         text_output += f"\nscheduler={args.scheduler}, device={args.device}"
-        text_output += (
-            f"\nsteps={args.steps}, guidance_scale={args.guidance_scale},"
-        )
-        text_output += (
-            f"seed={seeds[current_batch]}, size={args.height}x{args.width}"
-        )
+        text_output += f"\nsteps={args.steps}, guidance_scale={args.guidance_scale}, seed={seed}, size={args.height}x{args.width}"
         text_output += (
             f", batch size={args.batch_size}, max_length={args.max_length}"
         )

apps/stable_diffusion/scripts/txt2img_sdxl.py

@@ -1,96 +0,0 @@
-import torch
-import time
-from apps.stable_diffusion.src import (
-    args,
-    Text2ImageSDXLPipeline,
-    get_schedulers,
-    set_init_device_flags,
-    utils,
-    clear_all,
-    save_output_img,
-)
-
-
-def main():
-    if args.clear_all:
-        clear_all()
-
-    # TODO: prompt_embeds and text_embeds form base_model.json requires fixing
-    args.precision = "fp16"
-    args.height = 1024
-    args.width = 1024
-    args.max_length = 77
-    args.scheduler = "DDIM"
-    print(
-        "Using default supported configuration for SDXL :-\nprecision=fp16, width*height= 1024*1024, max_length=77 and scheduler=DDIM"
-    )
-    dtype = torch.float32 if args.precision == "fp32" else torch.half
-    cpu_scheduling = not args.scheduler.startswith("Shark")
-    set_init_device_flags()
-    schedulers = get_schedulers(args.hf_model_id)
-    scheduler_obj = schedulers[args.scheduler]
-    seed = args.seed
-    txt2img_obj = Text2ImageSDXLPipeline.from_pretrained(
-        scheduler=scheduler_obj,
-        import_mlir=args.import_mlir,
-        model_id=args.hf_model_id,
-        ckpt_loc=args.ckpt_loc,
-        precision=args.precision,
-        max_length=args.max_length,
-        batch_size=args.batch_size,
-        height=args.height,
-        width=args.width,
-        use_base_vae=args.use_base_vae,
-        use_tuned=args.use_tuned,
-        custom_vae=args.custom_vae,
-        low_cpu_mem_usage=args.low_cpu_mem_usage,
-        debug=args.import_debug if args.import_mlir else False,
-        use_lora=args.use_lora,
-        use_quantize=args.use_quantize,
-        ondemand=args.ondemand,
-    )
-
-    seeds = utils.batch_seeds(seed, args.batch_count, args.repeatable_seeds)
-    for current_batch in range(args.batch_count):
-        start_time = time.time()
-        generated_imgs = txt2img_obj.generate_images(
-            args.prompts,
-            args.negative_prompts,
-            args.batch_size,
-            args.height,
-            args.width,
-            args.steps,
-            args.guidance_scale,
-            seeds[current_batch],
-            args.max_length,
-            dtype,
-            args.use_base_vae,
-            cpu_scheduling,
-            args.max_embeddings_multiples,
-        )
-        total_time = time.time() - start_time
-        text_output = f"prompt={args.prompts}"
-        text_output += f"\nnegative prompt={args.negative_prompts}"
-        text_output += (
-            f"\nmodel_id={args.hf_model_id}, ckpt_loc={args.ckpt_loc}"
-        )
-        text_output += f"\nscheduler={args.scheduler}, device={args.device}"
-        text_output += (
-            f"\nsteps={args.steps}, guidance_scale={args.guidance_scale},"
-        )
-        text_output += (
-            f"seed={seeds[current_batch]}, size={args.height}x{args.width}"
-        )
-        text_output += (
-            f", batch size={args.batch_size}, max_length={args.max_length}"
-        )
-        # TODO: if using --batch_count=x txt2img_obj.log will output on each display every iteration infos from the start
-        text_output += txt2img_obj.log
-        text_output += f"\nTotal image generation time: {total_time:.4f}sec"
-
-        save_output_img(generated_imgs[0], seed)
-        print(text_output)
-
-
-if __name__ == "__main__":
-    main()

apps/stable_diffusion/scripts/upscaler.py

@@ -21,7 +21,7 @@ if __name__ == "__main__":
         print("Flag --img_path is required.")
         exit()
 
-    # When the models get uploaded, it should be defaulted to False.
+    # When the models get uploaded, it should be default to False.
     args.import_mlir = True
 
     cpu_scheduling = not args.scheduler.startswith("Shark")
@@ -73,7 +73,6 @@ if __name__ == "__main__":
         dtype,
         args.use_base_vae,
         cpu_scheduling,
-        args.max_embeddings_multiples,
     )
     total_time = time.time() - start_time
     text_output = f"prompt={args.prompts}"

apps/stable_diffusion/shark_sd.spec

@@ -1,13 +1,58 @@
 # -*- mode: python ; coding: utf-8 -*-
-from apps.stable_diffusion.shark_studio_imports import pathex, datas, hiddenimports
+from PyInstaller.utils.hooks import collect_data_files
+from PyInstaller.utils.hooks import copy_metadata
+from PyInstaller.utils.hooks import collect_submodules
+
+import sys ; sys.setrecursionlimit(sys.getrecursionlimit() * 5)
+
+datas = []
+datas += collect_data_files('torch')
+datas += copy_metadata('torch')
+datas += copy_metadata('tqdm')
+datas += copy_metadata('regex')
+datas += copy_metadata('requests')
+datas += copy_metadata('packaging')
+datas += copy_metadata('filelock')
+datas += copy_metadata('numpy')
+datas += copy_metadata('tokenizers')
+datas += copy_metadata('importlib_metadata')
+datas += copy_metadata('torch-mlir')
+datas += copy_metadata('omegaconf')
+datas += copy_metadata('safetensors')
+datas += copy_metadata('Pillow')
+datas += collect_data_files('diffusers')
+datas += collect_data_files('transformers')
+datas += collect_data_files('pytorch_lightning')
+datas += collect_data_files('opencv-python')
+datas += collect_data_files('skimage')
+datas += collect_data_files('gradio')
+datas += collect_data_files('gradio_client')
+datas += collect_data_files('iree')
+datas += collect_data_files('google-cloud-storage')
+datas += collect_data_files('shark')
+datas += collect_data_files('tkinter')
+datas += collect_data_files('webview')
+datas += collect_data_files('sentencepiece')
+datas += [
+         ( 'src/utils/resources/prompts.json', 'resources' ),
+         ( 'src/utils/resources/model_db.json', 'resources' ),
+         ( 'src/utils/resources/opt_flags.json', 'resources' ),
+         ( 'src/utils/resources/base_model.json', 'resources' ),
+         ( 'web/ui/css/*', 'ui/css' ),
+         ( 'web/ui/logos/*', 'logos' )
+         ]
 
 binaries = []
 
 block_cipher = None
 
+hiddenimports = ['shark', 'shark.shark_inference', 'apps']
+hiddenimports += [x for x in collect_submodules("skimage") if "tests" not in x]
+hiddenimports += [x for x in collect_submodules("iree") if "tests" not in x]
+
 a = Analysis(
     ['web/index.py'],
-    pathex=pathex,
+    pathex=['.'],
     binaries=binaries,
     datas=datas,
     hiddenimports=hiddenimports,
@@ -19,9 +64,6 @@ a = Analysis(
     win_private_assemblies=False,
     cipher=block_cipher,
     noarchive=False,
-    module_collection_mode={
-        'gradio': 'py',  # Collect gradio package as source .py files
-    },
 )
 pyz = PYZ(a.pure, a.zipped_data, cipher=block_cipher)
 
@@ -32,11 +74,11 @@ exe = EXE(
     a.zipfiles,
     a.datas,
     [],
-    name='nodai_shark_studio',
+    name='shark_sd',
     debug=False,
     bootloader_ignore_signals=False,
     strip=False,
-    upx=False,
+    upx=True,
     upx_exclude=[],
     runtime_tmpdir=None,
     console=True,

apps/stable_diffusion/shark_sd_cli.spec

@@ -29,7 +29,6 @@ datas += collect_data_files('gradio_client')
 datas += collect_data_files('iree')
 datas += collect_data_files('google-cloud-storage')
 datas += collect_data_files('shark')
-datas += collect_data_files('py-cpuinfo')
 datas += [
          ( 'src/utils/resources/prompts.json', 'resources' ),
          ( 'src/utils/resources/model_db.json', 'resources' ),

apps/stable_diffusion/shark_studio_imports.py

@@ -1,90 +0,0 @@
-from PyInstaller.utils.hooks import collect_data_files
-from PyInstaller.utils.hooks import copy_metadata
-from PyInstaller.utils.hooks import collect_submodules
-
-import sys
-
-sys.setrecursionlimit(sys.getrecursionlimit() * 5)
-
-# python path for pyinstaller
-pathex = [
-    ".",
-    "./apps/language_models/langchain",
-    "./apps/language_models/src/pipelines/minigpt4_utils",
-]
-
-# datafiles for pyinstaller
-datas = []
-datas += copy_metadata("torch")
-datas += copy_metadata("tokenizers")
-datas += copy_metadata("tqdm")
-datas += copy_metadata("regex")
-datas += copy_metadata("requests")
-datas += copy_metadata("packaging")
-datas += copy_metadata("filelock")
-datas += copy_metadata("numpy")
-datas += copy_metadata("importlib_metadata")
-datas += copy_metadata("torch-mlir")
-datas += copy_metadata("omegaconf")
-datas += copy_metadata("safetensors")
-datas += copy_metadata("Pillow")
-datas += copy_metadata("sentencepiece")
-datas += copy_metadata("pyyaml")
-datas += copy_metadata("huggingface-hub")
-datas += copy_metadata("gradio")
-datas += collect_data_files("torch")
-datas += collect_data_files("tokenizers")
-datas += collect_data_files("tiktoken")
-datas += collect_data_files("accelerate")
-datas += collect_data_files("diffusers")
-datas += collect_data_files("transformers")
-datas += collect_data_files("pytorch_lightning")
-datas += collect_data_files("skimage")
-datas += collect_data_files("gradio")
-datas += collect_data_files("gradio_client")
-datas += collect_data_files("iree")
-datas += collect_data_files("shark", include_py_files=True)
-datas += collect_data_files("timm", include_py_files=True)
-datas += collect_data_files("tqdm")
-datas += collect_data_files("tkinter")
-datas += collect_data_files("webview")
-datas += collect_data_files("sentencepiece")
-datas += collect_data_files("jsonschema")
-datas += collect_data_files("jsonschema_specifications")
-datas += collect_data_files("cpuinfo")
-datas += collect_data_files("langchain")
-datas += collect_data_files("cv2")
-datas += collect_data_files("einops")
-datas += [
-    ("src/utils/resources/prompts.json", "resources"),
-    ("src/utils/resources/model_db.json", "resources"),
-    ("src/utils/resources/opt_flags.json", "resources"),
-    ("src/utils/resources/base_model.json", "resources"),
-    ("web/ui/css/*", "ui/css"),
-    ("web/ui/logos/*", "logos"),
-    (
-        "../language_models/src/pipelines/minigpt4_utils/configs/*",
-        "minigpt4_utils/configs",
-    ),
-    (
-        "../language_models/src/pipelines/minigpt4_utils/prompts/*",
-        "minigpt4_utils/prompts",
-    ),
-]
-
-
-# hidden imports for pyinstaller
-hiddenimports = ["shark", "shark.shark_inference", "apps"]
-hiddenimports += [x for x in collect_submodules("skimage") if "tests" not in x]
-hiddenimports += [x for x in collect_submodules("gradio") if "tests" not in x]
-hiddenimports += [
-    x for x in collect_submodules("diffusers") if "tests" not in x
-]
-blacklist = ["tests", "convert"]
-hiddenimports += [
-    x
-    for x in collect_submodules("transformers")
-    if not any(kw in x for kw in blacklist)
-]
-hiddenimports += [x for x in collect_submodules("iree") if "tests" not in x]
-hiddenimports += ["iree._runtime"]

apps/stable_diffusion/src/__init__.py

@@ -9,7 +9,6 @@ from apps.stable_diffusion.src.utils import (
 )
 from apps.stable_diffusion.src.pipelines import (
     Text2ImagePipeline,
-    Text2ImageSDXLPipeline,
     Image2ImagePipeline,
     InpaintPipeline,
     OutpaintPipeline,

apps/stable_diffusion/src/models/opt_params.py

@@ -17,13 +17,9 @@ hf_model_variant_map = {
     "stabilityai/stable-diffusion-2-1-base": ["stablediffusion", "v2_1base"],
     "CompVis/stable-diffusion-v1-4": ["stablediffusion", "v1_4"],
     "runwayml/stable-diffusion-inpainting": ["stablediffusion", "inpaint_v1"],
-    "stabilityai/stable-diffusion-2-inpainting": [
-        "stablediffusion",
-        "inpaint_v2",
-    ],
+    "stabilityai/stable-diffusion-2-inpainting": ["stablediffusion", "inpaint_v2"],
 }
 
-
 # TODO: Add the quantized model as a part model_db.json.
 # This is currently in experimental phase.
 def get_quantize_model():
@@ -31,12 +27,9 @@ def get_quantize_model():
     model_key = "unet_int8"
     iree_flags = get_opt_flags("unet", precision="fp16")
     if args.height != 512 and args.width != 512 and args.max_length != 77:
-        sys.exit(
-            "The int8 quantized model currently requires the height and width to be 512, and max_length to be 77"
-        )
+        sys.exit("The int8 quantized model currently requires the height and width to be 512, and max_length to be 77")
     return bucket_key, model_key, iree_flags
 
-
 def get_variant_version(hf_model_id):
     return hf_model_variant_map[hf_model_id]
 
@@ -123,11 +116,8 @@ def get_clip():
     return get_shark_model(bucket, model_name, iree_flags)
 
 
-def get_tokenizer(subfolder="tokenizer", hf_model_id=None):
-    if hf_model_id is not None:
-        args.hf_model_id = hf_model_id
-
+def get_tokenizer():
     tokenizer = CLIPTokenizer.from_pretrained(
-        args.hf_model_id, subfolder=subfolder
+        args.hf_model_id, subfolder="tokenizer"
     )
     return tokenizer

apps/stable_diffusion/src/pipelines/__init__.py

@@ -1,9 +1,6 @@
 from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_txt2img import (
     Text2ImagePipeline,
 )
-from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_txt2img_sdxl import (
-    Text2ImageSDXLPipeline,
-)
 from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_img2img import (
     Image2ImagePipeline,
 )

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_img2img.py

@@ -15,11 +15,6 @@ from diffusers import (
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
     DEISMultistepScheduler,
-    DPMSolverSinglestepScheduler,
-    KDPM2AncestralDiscreteScheduler,
-    HeunDiscreteScheduler,
-    DDPMScheduler,
-    KDPM2DiscreteScheduler,
 )
 from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler
 from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import (
@@ -29,10 +24,6 @@ from apps.stable_diffusion.src.models import (
     SharkifyStableDiffusionModel,
     get_vae_encode,
 )
-from apps.stable_diffusion.src.utils import (
-    resamplers,
-    resampler_list,
-)
 
 
 class Image2ImagePipeline(StableDiffusionPipeline):
@@ -47,11 +38,6 @@ class Image2ImagePipeline(StableDiffusionPipeline):
             DPMSolverMultistepScheduler,
             SharkEulerDiscreteScheduler,
             DEISMultistepScheduler,
-            DPMSolverSinglestepScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            HeunDiscreteScheduler,
-            DDPMScheduler,
-            KDPM2DiscreteScheduler,
         ],
         sd_model: SharkifyStableDiffusionModel,
         import_mlir: bool,
@@ -88,21 +74,13 @@ class Image2ImagePipeline(StableDiffusionPipeline):
         num_inference_steps,
         strength,
         dtype,
-        resample_type,
     ):
         # Pre process image -> get image encoded -> process latents
 
         # TODO: process with variable HxW combos
 
-        # Pre-process image
-        resample_type = (
-            resamplers[resample_type]
-            if resample_type in resampler_list
-            # Fallback to Lanczos
-            else Image.Resampling.LANCZOS
-        )
-
-        image = image.resize((width, height), resample=resample_type)
+        # Pre process image
+        image = image.resize((width, height))
         image_arr = np.stack([np.array(i) for i in (image,)], axis=0)
         image_arr = image_arr / 255.0
         image_arr = torch.from_numpy(image_arr).permute(0, 3, 1, 2).to(dtype)
@@ -157,12 +135,7 @@ class Image2ImagePipeline(StableDiffusionPipeline):
         dtype,
         use_base_vae,
         cpu_scheduling,
-        max_embeddings_multiples,
-        stencils,
-        images,
-        resample_type,
-        control_mode,
-        preprocessed_hints=[],
+        use_stencil,
     ):
         # prompts and negative prompts must be a list.
         if isinstance(prompts, str):
@@ -183,10 +156,7 @@ class Image2ImagePipeline(StableDiffusionPipeline):
 
         # Get text embeddings with weight emphasis from prompts
         text_embeddings = self.encode_prompts_weight(
-            prompts,
-            neg_prompts,
-            max_length,
-            max_embeddings_multiples=max_embeddings_multiples,
+            prompts, neg_prompts, max_length
         )
 
         # guidance scale as a float32 tensor.
@@ -202,7 +172,6 @@ class Image2ImagePipeline(StableDiffusionPipeline):
             num_inference_steps=num_inference_steps,
             strength=strength,
             dtype=dtype,
-            resample_type=resample_type,
         )
 
         # Get Image latents

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_inpaint.py

@@ -14,11 +14,6 @@ from diffusers import (
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
     DEISMultistepScheduler,
-    DPMSolverSinglestepScheduler,
-    KDPM2AncestralDiscreteScheduler,
-    HeunDiscreteScheduler,
-    DDPMScheduler,
-    KDPM2DiscreteScheduler,
 )
 from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler
 from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import (
@@ -42,11 +37,6 @@ class InpaintPipeline(StableDiffusionPipeline):
             DPMSolverMultistepScheduler,
             SharkEulerDiscreteScheduler,
             DEISMultistepScheduler,
-            DPMSolverSinglestepScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            HeunDiscreteScheduler,
-            DDPMScheduler,
-            KDPM2DiscreteScheduler,
         ],
         sd_model: SharkifyStableDiffusionModel,
         import_mlir: bool,
@@ -388,7 +378,6 @@ class InpaintPipeline(StableDiffusionPipeline):
         dtype,
         use_base_vae,
         cpu_scheduling,
-        max_embeddings_multiples,
     ):
         # prompts and negative prompts must be a list.
         if isinstance(prompts, str):
@@ -419,10 +408,7 @@ class InpaintPipeline(StableDiffusionPipeline):
 
         # Get text embeddings with weight emphasis from prompts
         text_embeddings = self.encode_prompts_weight(
-            prompts,
-            neg_prompts,
-            max_length,
-            max_embeddings_multiples=max_embeddings_multiples,
+            prompts, neg_prompts, max_length
         )
 
         # guidance scale as a float32 tensor.

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_outpaint.py

@@ -14,11 +14,6 @@ from diffusers import (
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
     DEISMultistepScheduler,
-    DPMSolverSinglestepScheduler,
-    KDPM2AncestralDiscreteScheduler,
-    HeunDiscreteScheduler,
-    DDPMScheduler,
-    KDPM2DiscreteScheduler,
 )
 from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler
 from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import (
@@ -43,11 +38,6 @@ class OutpaintPipeline(StableDiffusionPipeline):
             DPMSolverMultistepScheduler,
             SharkEulerDiscreteScheduler,
             DEISMultistepScheduler,
-            DPMSolverSinglestepScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            HeunDiscreteScheduler,
-            DDPMScheduler,
-            KDPM2DiscreteScheduler,
         ],
         sd_model: SharkifyStableDiffusionModel,
         import_mlir: bool,
@@ -389,7 +379,6 @@ class OutpaintPipeline(StableDiffusionPipeline):
         dtype,
         use_base_vae,
         cpu_scheduling,
-        max_embeddings_multiples,
     ):
         # prompts and negative prompts must be a list.
         if isinstance(prompts, str):
@@ -420,10 +409,7 @@ class OutpaintPipeline(StableDiffusionPipeline):
 
         # Get text embeddings with weight emphasis from prompts
         text_embeddings = self.encode_prompts_weight(
-            prompts,
-            neg_prompts,
-            max_length,
-            max_embeddings_multiples=max_embeddings_multiples,
+            prompts, neg_prompts, max_length
         )
 
         # guidance scale as a float32 tensor.

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_stencil.py

@@ -14,33 +14,17 @@ from diffusers import (
     EulerDiscreteScheduler,
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
-    DEISMultistepScheduler,
-    DPMSolverSinglestepScheduler,
-    KDPM2AncestralDiscreteScheduler,
-    HeunDiscreteScheduler,
-    DDPMScheduler,
-    KDPM2DiscreteScheduler,
 )
 from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler
 from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import (
     StableDiffusionPipeline,
 )
-from apps.stable_diffusion.src.utils import (
-    controlnet_hint_conversion,
-    controlnet_hint_reshaping,
-)
+from apps.stable_diffusion.src.utils import controlnet_hint_conversion
 from apps.stable_diffusion.src.utils import (
     start_profiling,
     end_profiling,
 )
-from apps.stable_diffusion.src.utils import (
-    resamplers,
-    resampler_list,
-)
-from apps.stable_diffusion.src.models import (
-    SharkifyStableDiffusionModel,
-    get_vae_encode,
-)
+from apps.stable_diffusion.src.models import SharkifyStableDiffusionModel
 
 
 class StencilPipeline(StableDiffusionPipeline):
@@ -54,76 +38,23 @@ class StencilPipeline(StableDiffusionPipeline):
             EulerAncestralDiscreteScheduler,
             DPMSolverMultistepScheduler,
             SharkEulerDiscreteScheduler,
-            DEISMultistepScheduler,
-            DPMSolverSinglestepScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            HeunDiscreteScheduler,
-            DDPMScheduler,
-            KDPM2DiscreteScheduler,
         ],
         sd_model: SharkifyStableDiffusionModel,
         import_mlir: bool,
         use_lora: str,
         ondemand: bool,
-        controlnet_names: list[str],
     ):
         super().__init__(scheduler, sd_model, import_mlir, use_lora, ondemand)
-        self.controlnet = [None] * len(controlnet_names)
-        self.controlnet_512 = [None] * len(controlnet_names)
-        self.controlnet_id = [str] * len(controlnet_names)
-        self.controlnet_512_id = [str] * len(controlnet_names)
-        self.controlnet_names = controlnet_names
-        self.vae_encode = None
+        self.controlnet = None
 
-    def load_vae_encode(self):
-        if self.vae_encode is not None:
+    def load_controlnet(self):
+        if self.controlnet is not None:
             return
+        self.controlnet = self.sd_model.controlnet()
 
-        if self.import_mlir or self.use_lora:
-            self.vae_encode = self.sd_model.vae_encode()
-        else:
-            try:
-                self.vae_encode = get_vae_encode()
-            except:
-                print("download pipeline failed, falling back to import_mlir")
-                self.vae_encode = self.sd_model.vae_encode()
-
-    def unload_vae_encode(self):
-        del self.vae_encode
-        self.vae_encode = None
-
-    def load_controlnet(self, index, model_name):
-        if model_name is None:
-            return
-        if (
-            self.controlnet[index] is not None
-            and self.controlnet_id[index] is not None
-            and self.controlnet_id[index] == model_name
-        ):
-            return
-        self.controlnet_id[index] = model_name
-        self.controlnet[index] = self.sd_model.controlnet(model_name)
-
-    def unload_controlnet(self, index):
-        del self.controlnet[index]
-        self.controlnet_id[index] = None
-        self.controlnet[index] = None
-
-    def load_controlnet_512(self, index, model_name):
-        if (
-            self.controlnet_512[index] is not None
-            and self.controlnet_512_id[index] == model_name
-        ):
-            return
-        self.controlnet_512_id[index] = model_name
-        self.controlnet_512[index] = self.sd_model.controlnet(
-            model_name, use_large=True
-        )
-
-    def unload_controlnet_512(self, index):
-        del self.controlnet_512[index]
-        self.controlnet_512_id[index] = None
-        self.controlnet_512[index] = None
+    def unload_controlnet(self):
+        del self.controlnet
+        self.controlnet = None
 
     def prepare_latents(
         self,
@@ -150,58 +81,6 @@ class StencilPipeline(StableDiffusionPipeline):
         latents = latents * self.scheduler.init_noise_sigma
         return latents
 
-    def prepare_image_latents(
-        self,
-        image,
-        batch_size,
-        height,
-        width,
-        generator,
-        num_inference_steps,
-        strength,
-        dtype,
-        resample_type,
-    ):
-        # Pre process image -> get image encoded -> process latents
-
-        # TODO: process with variable HxW combos
-
-        # Pre-process image
-        resample_type = (
-            resamplers[resample_type]
-            if resample_type in resampler_list
-            # Fallback to Lanczos
-            else Image.Resampling.LANCZOS
-        )
-
-        image = image.resize((width, height), resample=resample_type)
-        image_arr = np.stack([np.array(i) for i in (image,)], axis=0)
-        image_arr = image_arr / 255.0
-        image_arr = torch.from_numpy(image_arr).permute(0, 3, 1, 2).to(dtype)
-        image_arr = 2 * (image_arr - 0.5)
-
-        # set scheduler steps
-        self.scheduler.set_timesteps(num_inference_steps)
-        init_timestep = min(
-            int(num_inference_steps * strength), num_inference_steps
-        )
-        t_start = max(num_inference_steps - init_timestep, 0)
-        # timesteps reduced as per strength
-        timesteps = self.scheduler.timesteps[t_start:]
-        # new number of steps to be used as per strength will be
-        # num_inference_steps = num_inference_steps - t_start
-
-        # image encode
-        latents = self.encode_image((image_arr,))
-        latents = torch.from_numpy(latents).to(dtype)
-        # add noise to data
-        noise = torch.randn(latents.shape, generator=generator, dtype=dtype)
-        latents = self.scheduler.add_noise(
-            latents, noise, timesteps[0].repeat(1)
-        )
-
-        return latents, timesteps
-
     def produce_stencil_latents(
         self,
         latents,
@@ -210,9 +89,8 @@ class StencilPipeline(StableDiffusionPipeline):
         total_timesteps,
         dtype,
         cpu_scheduling,
-        stencil_hints=[None],
+        controlnet_hint=None,
         controlnet_conditioning_scale: float = 1.0,
-        control_mode="Balanced",  # Prompt, Balanced, or Controlnet
         mask=None,
         masked_image_latents=None,
         return_all_latents=False,
@@ -221,24 +99,8 @@ class StencilPipeline(StableDiffusionPipeline):
         latent_history = [latents]
         text_embeddings = torch.from_numpy(text_embeddings).to(dtype)
         text_embeddings_numpy = text_embeddings.detach().numpy()
-        assert control_mode in ["Prompt", "Balanced", "Controlnet"]
-        if text_embeddings.shape[1] <= self.model_max_length:
-            self.load_unet()
-        else:
-            self.load_unet_512()
-
-        for i, name in enumerate(self.controlnet_names):
-            use_names = []
-            if name is not None:
-                use_names.append(name)
-            else:
-                continue
-            if text_embeddings.shape[1] <= self.model_max_length:
-                self.load_controlnet(i, name)
-            else:
-                self.load_controlnet_512(i, name)
-            self.controlnet_names = use_names
-
+        self.load_unet()
+        self.load_controlnet()
         for i, t in tqdm(enumerate(total_timesteps)):
             step_start_time = time.time()
             timestep = torch.tensor([t]).to(dtype)
@@ -261,167 +123,43 @@ class StencilPipeline(StableDiffusionPipeline):
                 ).to(dtype)
             else:
                 latent_model_input_1 = latent_model_input
-
-            # Multicontrolnet
-            width = latent_model_input_1.shape[2]
-            height = latent_model_input_1.shape[3]
-            dtype = latent_model_input_1.dtype
-            control_acc = (
-                [torch.zeros((2, 320, height, width), dtype=dtype)] * 3
-                + [
-                    torch.zeros(
-                        (2, 320, int(height / 2), int(width / 2)), dtype=dtype
-                    )
-                ]
-                + [
-                    torch.zeros(
-                        (2, 640, int(height / 2), int(width / 2)), dtype=dtype
-                    )
-                ]
-                * 2
-                + [
-                    torch.zeros(
-                        (2, 640, int(height / 4), int(width / 4)), dtype=dtype
-                    )
-                ]
-                + [
-                    torch.zeros(
-                        (2, 1280, int(height / 4), int(width / 4)), dtype=dtype
-                    )
-                ]
-                * 2
-                + [
-                    torch.zeros(
-                        (2, 1280, int(height / 8), int(width / 8)), dtype=dtype
-                    )
-                ]
-                * 4
+            control = self.controlnet(
+                "forward",
+                (
+                    latent_model_input_1,
+                    timestep,
+                    text_embeddings,
+                    controlnet_hint,
+                ),
+                send_to_host=False,
             )
-            for i, controlnet_hint in enumerate(stencil_hints):
-                if controlnet_hint is None:
-                    pass
-                if text_embeddings.shape[1] <= self.model_max_length:
-                    control = self.controlnet[i](
-                        "forward",
-                        (
-                            latent_model_input_1,
-                            timestep,
-                            text_embeddings,
-                            controlnet_hint,
-                            *control_acc,
-                        ),
-                        send_to_host=False,
-                    )
-                else:
-                    control = self.controlnet_512[i](
-                        "forward",
-                        (
-                            latent_model_input_1,
-                            timestep,
-                            text_embeddings,
-                            controlnet_hint,
-                            *control_acc,
-                        ),
-                        send_to_host=False,
-                    )
-                control_acc = control[13:]
-                control = control[:13]
-
             timestep = timestep.detach().numpy()
             # Profiling Unet.
             profile_device = start_profiling(file_path="unet.rdc")
             # TODO: Pass `control` as it is to Unet. Same as TODO mentioned in model_wrappers.py.
-
-            dtype = latents.dtype
-            if control_mode == "Balanced":
-                control_scale = [
-                    torch.tensor(1.0, dtype=dtype) for _ in range(len(control))
-                ]
-            elif control_mode == "Prompt":
-                control_scale = [
-                    torch.tensor(0.825**x, dtype=dtype)
-                    for x in range(len(control))
-                ]
-            elif control_mode == "Controlnet":
-                control_scale = [
-                    torch.tensor(float(guidance_scale), dtype=dtype)
-                    for _ in range(len(control))
-                ]
-
-            if text_embeddings.shape[1] <= self.model_max_length:
-                noise_pred = self.unet(
-                    "forward",
-                    (
-                        latent_model_input,
-                        timestep,
-                        text_embeddings_numpy,
-                        guidance_scale,
-                        control[0],
-                        control[1],
-                        control[2],
-                        control[3],
-                        control[4],
-                        control[5],
-                        control[6],
-                        control[7],
-                        control[8],
-                        control[9],
-                        control[10],
-                        control[11],
-                        control[12],
-                        control_scale[0],
-                        control_scale[1],
-                        control_scale[2],
-                        control_scale[3],
-                        control_scale[4],
-                        control_scale[5],
-                        control_scale[6],
-                        control_scale[7],
-                        control_scale[8],
-                        control_scale[9],
-                        control_scale[10],
-                        control_scale[11],
-                        control_scale[12],
-                    ),
-                    send_to_host=False,
-                )
-            else:
-                noise_pred = self.unet_512(
-                    "forward",
-                    (
-                        latent_model_input,
-                        timestep,
-                        text_embeddings_numpy,
-                        guidance_scale,
-                        control[0],
-                        control[1],
-                        control[2],
-                        control[3],
-                        control[4],
-                        control[5],
-                        control[6],
-                        control[7],
-                        control[8],
-                        control[9],
-                        control[10],
-                        control[11],
-                        control[12],
-                        control_scale[0],
-                        control_scale[1],
-                        control_scale[2],
-                        control_scale[3],
-                        control_scale[4],
-                        control_scale[5],
-                        control_scale[6],
-                        control_scale[7],
-                        control_scale[8],
-                        control_scale[9],
-                        control_scale[10],
-                        control_scale[11],
-                        control_scale[12],
-                    ),
-                    send_to_host=False,
-                )
+            noise_pred = self.unet(
+                "forward",
+                (
+                    latent_model_input,
+                    timestep,
+                    text_embeddings_numpy,
+                    guidance_scale,
+                    control[0],
+                    control[1],
+                    control[2],
+                    control[3],
+                    control[4],
+                    control[5],
+                    control[6],
+                    control[7],
+                    control[8],
+                    control[9],
+                    control[10],
+                    control[11],
+                    control[12],
+                ),
+                send_to_host=False,
+            )
             end_profiling(profile_device)
 
             if cpu_scheduling:
@@ -441,10 +179,7 @@ class StencilPipeline(StableDiffusionPipeline):
 
         if self.ondemand:
             self.unload_unet()
-            self.unload_unet_512()
-            for i in range(len(self.controlnet_names)):
-                self.unload_controlnet(i)
-                self.unload_controlnet_512(i)
+            self.unload_controlnet()
         avg_step_time = step_time_sum / len(total_timesteps)
         self.log += f"\nAverage step time: {avg_step_time}ms/it"
 
@@ -453,17 +188,6 @@ class StencilPipeline(StableDiffusionPipeline):
         all_latents = torch.cat(latent_history, dim=0)
         return all_latents
 
-    def encode_image(self, input_image):
-        self.load_vae_encode()
-        vae_encode_start = time.time()
-        latents = self.vae_encode("forward", input_image)
-        vae_inf_time = (time.time() - vae_encode_start) * 1000
-        if self.ondemand:
-            self.unload_vae_encode()
-        self.log += f"\nVAE Encode Inference time (ms): {vae_inf_time:.3f}"
-
-        return latents
-
     def generate_images(
         self,
         prompts,
@@ -480,49 +204,13 @@ class StencilPipeline(StableDiffusionPipeline):
         dtype,
         use_base_vae,
         cpu_scheduling,
-        max_embeddings_multiples,
-        stencils,
-        stencil_images,
-        resample_type,
-        control_mode,
-        preprocessed_hints,
+        use_stencil,
     ):
         # Control Embedding check & conversion
-        # controlnet_hint = controlnet_hint_conversion(
-        #     image, use_stencil, height, width, dtype, num_images_per_prompt=1
-        # )
-        stencil_hints = []
-        self.sd_model.stencils = stencils
-        for i, hint in enumerate(preprocessed_hints):
-            if hint is not None:
-                hint = controlnet_hint_reshaping(
-                    hint,
-                    height,
-                    width,
-                    dtype,
-                    num_images_per_prompt=1,
-                )
-                stencil_hints.append(hint)
-
-        for i, stencil in enumerate(stencils):
-            if stencil == None:
-                continue
-            if len(stencil_hints) > i:
-                if stencil_hints[i] is not None:
-                    print(f"Using preprocessed controlnet hint for {stencil}")
-                    continue
-            image = stencil_images[i]
-            stencil_hints.append(
-                controlnet_hint_conversion(
-                    image,
-                    stencil,
-                    height,
-                    width,
-                    dtype,
-                    num_images_per_prompt=1,
-                )
-            )
-
+        # TODO: 1. Change `num_images_per_prompt`.
+        controlnet_hint = controlnet_hint_conversion(
+            image, use_stencil, height, width, dtype, num_images_per_prompt=1
+        )
         # prompts and negative prompts must be a list.
         if isinstance(prompts, str):
             prompts = [prompts]
@@ -542,38 +230,22 @@ class StencilPipeline(StableDiffusionPipeline):
 
         # Get text embeddings with weight emphasis from prompts
         text_embeddings = self.encode_prompts_weight(
-            prompts,
-            neg_prompts,
-            max_length,
-            max_embeddings_multiples=max_embeddings_multiples,
+            prompts, neg_prompts, max_length
         )
 
         # guidance scale as a float32 tensor.
         guidance_scale = torch.tensor(guidance_scale).to(torch.float32)
-        if image is not None:
-            # Prepare input image latent
-            init_latents, final_timesteps = self.prepare_image_latents(
-                image=image,
-                batch_size=batch_size,
-                height=height,
-                width=width,
-                generator=generator,
-                num_inference_steps=num_inference_steps,
-                strength=strength,
-                dtype=dtype,
-                resample_type=resample_type,
-            )
-        else:
-            # Prepare initial latent.
-            init_latents = self.prepare_latents(
-                batch_size=batch_size,
-                height=height,
-                width=width,
-                generator=generator,
-                num_inference_steps=num_inference_steps,
-                dtype=dtype,
-            )
-            final_timesteps = self.scheduler.timesteps
+
+        # Prepare initial latent.
+        init_latents = self.prepare_latents(
+            batch_size=batch_size,
+            height=height,
+            width=width,
+            generator=generator,
+            num_inference_steps=num_inference_steps,
+            dtype=dtype,
+        )
+        final_timesteps = self.scheduler.timesteps
 
         # Get Image latents
         latents = self.produce_stencil_latents(
@@ -583,8 +255,7 @@ class StencilPipeline(StableDiffusionPipeline):
             total_timesteps=final_timesteps,
             dtype=dtype,
             cpu_scheduling=cpu_scheduling,
-            control_mode=control_mode,
-            stencil_hints=stencil_hints,
+            controlnet_hint=controlnet_hint,
         )
 
         # Img latents -> PIL images

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_txt2img.py

@@ -13,15 +13,8 @@ from diffusers import (
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
     DEISMultistepScheduler,
-    DDPMScheduler,
-    DPMSolverSinglestepScheduler,
-    KDPM2AncestralDiscreteScheduler,
-    HeunDiscreteScheduler,
-)
-from apps.stable_diffusion.src.schedulers import (
-    SharkEulerDiscreteScheduler,
-    SharkEulerAncestralDiscreteScheduler,
 )
+from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler
 from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import (
     StableDiffusionPipeline,
 )
@@ -41,10 +34,6 @@ class Text2ImagePipeline(StableDiffusionPipeline):
             DPMSolverMultistepScheduler,
             SharkEulerDiscreteScheduler,
             DEISMultistepScheduler,
-            DDPMScheduler,
-            DPMSolverSinglestepScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            HeunDiscreteScheduler,
         ],
         sd_model: SharkifyStableDiffusionModel,
         import_mlir: bool,

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_txt2img_sdxl.py

@@ -1,220 +0,0 @@
-import torch
-import numpy as np
-from random import randint
-from typing import Union
-from diffusers import (
-    DDIMScheduler,
-    PNDMScheduler,
-    LMSDiscreteScheduler,
-    KDPM2DiscreteScheduler,
-    EulerDiscreteScheduler,
-    EulerAncestralDiscreteScheduler,
-    DPMSolverMultistepScheduler,
-    DEISMultistepScheduler,
-    DDPMScheduler,
-    DPMSolverSinglestepScheduler,
-    KDPM2AncestralDiscreteScheduler,
-    HeunDiscreteScheduler,
-)
-from apps.stable_diffusion.src.schedulers import (
-    SharkEulerDiscreteScheduler,
-    SharkEulerAncestralDiscreteScheduler,
-)
-from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import (
-    StableDiffusionPipeline,
-)
-from apps.stable_diffusion.src.models import SharkifyStableDiffusionModel
-from transformers.utils import logging
-
-logger = logging.get_logger(__name__)
-
-
-class Text2ImageSDXLPipeline(StableDiffusionPipeline):
-    def __init__(
-        self,
-        scheduler: Union[
-            DDIMScheduler,
-            PNDMScheduler,
-            LMSDiscreteScheduler,
-            KDPM2DiscreteScheduler,
-            EulerDiscreteScheduler,
-            EulerAncestralDiscreteScheduler,
-            DPMSolverMultistepScheduler,
-            SharkEulerDiscreteScheduler,
-            SharkEulerAncestralDiscreteScheduler,
-            DEISMultistepScheduler,
-            DDPMScheduler,
-            DPMSolverSinglestepScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            HeunDiscreteScheduler,
-        ],
-        sd_model: SharkifyStableDiffusionModel,
-        import_mlir: bool,
-        use_lora: str,
-        ondemand: bool,
-        is_fp32_vae: bool,
-    ):
-        super().__init__(scheduler, sd_model, import_mlir, use_lora, ondemand)
-        self.is_fp32_vae = is_fp32_vae
-
-    def prepare_latents(
-        self,
-        batch_size,
-        height,
-        width,
-        generator,
-        num_inference_steps,
-        dtype,
-    ):
-        latents = torch.randn(
-            (
-                batch_size,
-                4,
-                height // 8,
-                width // 8,
-            ),
-            generator=generator,
-            dtype=torch.float32,
-        ).to(dtype)
-
-        self.scheduler.set_timesteps(num_inference_steps)
-        self.scheduler.is_scale_input_called = True
-        latents = latents * self.scheduler.init_noise_sigma
-        return latents
-
-    def _get_add_time_ids(
-        self, original_size, crops_coords_top_left, target_size, dtype
-    ):
-        add_time_ids = list(
-            original_size + crops_coords_top_left + target_size
-        )
-
-        # self.unet.config.addition_time_embed_dim IS 256.
-        # self.text_encoder_2.config.projection_dim IS 1280.
-        passed_add_embed_dim = 256 * len(add_time_ids) + 1280
-        expected_add_embed_dim = 2816
-        # self.unet.add_embedding.linear_1.in_features IS 2816.
-
-        if expected_add_embed_dim != passed_add_embed_dim:
-            raise ValueError(
-                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
-            )
-
-        add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
-        return add_time_ids
-
-    def generate_images(
-        self,
-        prompts,
-        neg_prompts,
-        batch_size,
-        height,
-        width,
-        num_inference_steps,
-        guidance_scale,
-        seed,
-        max_length,
-        dtype,
-        use_base_vae,
-        cpu_scheduling,
-        max_embeddings_multiples,
-    ):
-        # prompts and negative prompts must be a list.
-        if isinstance(prompts, str):
-            prompts = [prompts]
-
-        if isinstance(neg_prompts, str):
-            neg_prompts = [neg_prompts]
-
-        prompts = prompts * batch_size
-        neg_prompts = neg_prompts * batch_size
-
-        # seed generator to create the inital latent noise. Also handle out of range seeds.
-        # TODO: Wouldn't it be preferable to just report an error instead of modifying the seed on the fly?
-        uint32_info = np.iinfo(np.uint32)
-        uint32_min, uint32_max = uint32_info.min, uint32_info.max
-        if seed < uint32_min or seed >= uint32_max:
-            seed = randint(uint32_min, uint32_max)
-        generator = torch.manual_seed(seed)
-
-        # Get initial latents.
-        init_latents = self.prepare_latents(
-            batch_size=batch_size,
-            height=height,
-            width=width,
-            generator=generator,
-            num_inference_steps=num_inference_steps,
-            dtype=dtype,
-        )
-
-        # Get text embeddings.
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = self.encode_prompt_sdxl(
-            prompt=prompts,
-            num_images_per_prompt=1,
-            do_classifier_free_guidance=True,
-            negative_prompt=neg_prompts,
-        )
-
-        # Prepare timesteps.
-        self.scheduler.set_timesteps(num_inference_steps)
-
-        timesteps = self.scheduler.timesteps
-
-        # Prepare added time ids & embeddings.
-        original_size = (height, width)
-        target_size = (height, width)
-        crops_coords_top_left = (0, 0)
-        add_text_embeds = pooled_prompt_embeds
-        add_time_ids = self._get_add_time_ids(
-            original_size,
-            crops_coords_top_left,
-            target_size,
-            dtype=prompt_embeds.dtype,
-        )
-
-        prompt_embeds = torch.cat(
-            [negative_prompt_embeds, prompt_embeds], dim=0
-        )
-        add_text_embeds = torch.cat(
-            [negative_pooled_prompt_embeds, add_text_embeds], dim=0
-        )
-        add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
-
-        prompt_embeds = prompt_embeds
-        add_text_embeds = add_text_embeds.to(dtype)
-        add_time_ids = add_time_ids.repeat(batch_size * 1, 1)
-
-        # guidance scale as a float32 tensor.
-        guidance_scale = torch.tensor(guidance_scale).to(dtype)
-        prompt_embeds = prompt_embeds.to(dtype)
-        add_time_ids = add_time_ids.to(dtype)
-
-        # Get Image latents.
-        latents = self.produce_img_latents_sdxl(
-            init_latents,
-            timesteps,
-            add_text_embeds,
-            add_time_ids,
-            prompt_embeds,
-            cpu_scheduling,
-            guidance_scale,
-            dtype,
-        )
-
-        # Img latents -> PIL images.
-        all_imgs = []
-        self.load_vae()
-        for i in range(0, latents.shape[0], batch_size):
-            imgs = self.decode_latents_sdxl(
-                latents[i : i + batch_size], is_fp32_vae=self.is_fp32_vae
-            )
-            all_imgs.extend(imgs)
-        if self.ondemand:
-            self.unload_vae()
-
-        return all_imgs

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_upscaler.py

@@ -17,9 +17,6 @@ from diffusers import (
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
     DEISMultistepScheduler,
-    DPMSolverSinglestepScheduler,
-    KDPM2AncestralDiscreteScheduler,
-    HeunDiscreteScheduler,
 )
 from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler
 from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import (
@@ -70,11 +67,6 @@ class UpscalerPipeline(StableDiffusionPipeline):
             DPMSolverMultistepScheduler,
             SharkEulerDiscreteScheduler,
             DEISMultistepScheduler,
-            DDPMScheduler,
-            DPMSolverSinglestepScheduler,
-            KDPM2DiscreteScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            HeunDiscreteScheduler,
         ],
         low_res_scheduler: Union[
             DDIMScheduler,
@@ -86,10 +78,6 @@ class UpscalerPipeline(StableDiffusionPipeline):
             DPMSolverMultistepScheduler,
             SharkEulerDiscreteScheduler,
             DEISMultistepScheduler,
-            DPMSolverSinglestepScheduler,
-            KDPM2DiscreteScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            HeunDiscreteScheduler,
         ],
         sd_model: SharkifyStableDiffusionModel,
         import_mlir: bool,
@@ -180,10 +168,7 @@ class UpscalerPipeline(StableDiffusionPipeline):
         text_embeddings = torch.from_numpy(text_embeddings).to(dtype)
         text_embeddings_numpy = text_embeddings.detach().numpy()
         self.status = SD_STATE_IDLE
-        if text_embeddings.shape[1] <= self.model_max_length:
-            self.load_unet()
-        else:
-            self.load_unet_512()
+        self.load_unet()
         for i, t in tqdm(enumerate(total_timesteps)):
             step_start_time = time.time()
             latent_model_input = torch.cat([latents] * 2)
@@ -197,26 +182,15 @@ class UpscalerPipeline(StableDiffusionPipeline):
 
             # Profiling Unet.
             profile_device = start_profiling(file_path="unet.rdc")
-            if text_embeddings.shape[1] <= self.model_max_length:
-                noise_pred = self.unet(
-                    "forward",
-                    (
-                        latent_model_input,
-                        timestep,
-                        text_embeddings_numpy,
-                        noise_level,
-                    ),
-                )
-            else:
-                noise_pred = self.unet_512(
-                    "forward",
-                    (
-                        latent_model_input,
-                        timestep,
-                        text_embeddings_numpy,
-                        noise_level,
-                    ),
-                )
+            noise_pred = self.unet(
+                "forward",
+                (
+                    latent_model_input,
+                    timestep,
+                    text_embeddings_numpy,
+                    noise_level,
+                ),
+            )
             end_profiling(profile_device)
             noise_pred = torch.from_numpy(noise_pred)
             noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
@@ -245,7 +219,6 @@ class UpscalerPipeline(StableDiffusionPipeline):
 
         if self.ondemand:
             self.unload_unet()
-            self.unload_unet_512()
         avg_step_time = step_time_sum / len(total_timesteps)
         self.log += f"\nAverage step time: {avg_step_time}ms/it"
 
@@ -270,7 +243,6 @@ class UpscalerPipeline(StableDiffusionPipeline):
         dtype,
         use_base_vae,
         cpu_scheduling,
-        max_embeddings_multiples,
     ):
         # prompts and negative prompts must be a list.
         if isinstance(prompts, str):
@@ -292,10 +264,7 @@ class UpscalerPipeline(StableDiffusionPipeline):
 
         # Get text embeddings with weight emphasis from prompts
         text_embeddings = self.encode_prompts_weight(
-            prompts,
-            neg_prompts,
-            max_length,
-            max_embeddings_multiples=max_embeddings_multiples,
+            prompts, neg_prompts, max_length
         )
 
         # 4. Preprocess image

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_utils.py

@@ -15,15 +15,9 @@ from diffusers import (
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
     DEISMultistepScheduler,
-    DPMSolverSinglestepScheduler,
-    KDPM2AncestralDiscreteScheduler,
-    HeunDiscreteScheduler,
 )
 from shark.shark_inference import SharkInference
-from apps.stable_diffusion.src.schedulers import (
-    SharkEulerDiscreteScheduler,
-    SharkEulerAncestralDiscreteScheduler,
-)
+from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler
 from apps.stable_diffusion.src.models import (
     SharkifyStableDiffusionModel,
     get_vae,
@@ -36,8 +30,6 @@ from apps.stable_diffusion.src.utils import (
     end_profiling,
 )
 import sys
-import gc
-from typing import List, Optional
 
 SD_STATE_IDLE = "idle"
 SD_STATE_CANCEL = "cancel"
@@ -55,36 +47,27 @@ class StableDiffusionPipeline:
             EulerAncestralDiscreteScheduler,
             DPMSolverMultistepScheduler,
             SharkEulerDiscreteScheduler,
-            SharkEulerAncestralDiscreteScheduler,
             DEISMultistepScheduler,
-            DDPMScheduler,
-            DPMSolverSinglestepScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            HeunDiscreteScheduler,
         ],
         sd_model: SharkifyStableDiffusionModel,
         import_mlir: bool,
         use_lora: str,
         ondemand: bool,
-        is_f32_vae: bool = False,
     ):
         self.vae = None
         self.text_encoder = None
-        self.text_encoder_2 = None
         self.unet = None
         self.unet_512 = None
         self.model_max_length = 77
+        self.scheduler = scheduler
         # TODO: Implement using logging python utility.
         self.log = ""
         self.status = SD_STATE_IDLE
         self.sd_model = sd_model
-        self.scheduler = scheduler
         self.import_mlir = import_mlir
         self.use_lora = use_lora
         self.ondemand = ondemand
-        self.is_f32_vae = is_f32_vae
-        # TODO: Find a better workaround for fetching base_model_id early
-        #  enough for CLIPTokenizer.
+        # TODO: Find a better workaround for fetching base_model_id early enough for CLIPTokenizer.
         try:
             self.tokenizer = get_tokenizer()
         except:
@@ -99,8 +82,7 @@ class StableDiffusionPipeline:
         if self.import_mlir or self.use_lora:
             if not self.import_mlir:
                 print(
-                    "Warning: LoRA provided but import_mlir not specified. "
-                    "Importing MLIR anyways."
+                    "Warning: LoRA provided but import_mlir not specified. Importing MLIR anyways."
                 )
             self.text_encoder = self.sd_model.clip()
         else:
@@ -115,34 +97,6 @@ class StableDiffusionPipeline:
         del self.text_encoder
         self.text_encoder = None
 
-    def load_clip_sdxl(self):
-        if self.text_encoder and self.text_encoder_2:
-            return
-
-        if self.import_mlir or self.use_lora:
-            if not self.import_mlir:
-                print(
-                    "Warning: LoRA provided but import_mlir not specified. "
-                    "Importing MLIR anyways."
-                )
-            self.text_encoder, self.text_encoder_2 = self.sd_model.sdxl_clip()
-        else:
-            try:
-                # TODO: Fix this for SDXL
-                self.text_encoder = get_clip()
-            except Exception as e:
-                print(e)
-                print("download pipeline failed, falling back to import_mlir")
-                (
-                    self.text_encoder,
-                    self.text_encoder_2,
-                ) = self.sd_model.sdxl_clip()
-
-    def unload_clip_sdxl(self):
-        del self.text_encoder, self.text_encoder_2
-        self.text_encoder = None
-        self.text_encoder_2 = None
-
     def load_unet(self):
         if self.unet is not None:
             return
@@ -196,182 +150,6 @@ class StableDiffusionPipeline:
     def unload_vae(self):
         del self.vae
         self.vae = None
-        gc.collect()
-
-    def encode_prompt_sdxl(
-        self,
-        prompt: str,
-        num_images_per_prompt: int = 1,
-        do_classifier_free_guidance: bool = True,
-        negative_prompt: Optional[str] = None,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
-        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        hf_model_id: Optional[
-            str
-        ] = "stabilityai/stable-diffusion-xl-base-1.0",
-    ):
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        # Define tokenizers and text encoders
-        self.tokenizer_2 = get_tokenizer("tokenizer_2", hf_model_id)
-        self.load_clip_sdxl()
-        tokenizers = (
-            [self.tokenizer, self.tokenizer_2]
-            if self.tokenizer is not None
-            else [self.tokenizer_2]
-        )
-        text_encoders = (
-            [self.text_encoder, self.text_encoder_2]
-            if self.text_encoder is not None
-            else [self.text_encoder_2]
-        )
-
-        # textual inversion: procecss multi-vector tokens if necessary
-        prompt_embeds_list = []
-        prompts = [prompt, prompt]
-        for prompt, tokenizer, text_encoder in zip(
-            prompts, tokenizers, text_encoders
-        ):
-            text_inputs = tokenizer(
-                prompt,
-                padding="max_length",
-                max_length=tokenizer.model_max_length,
-                truncation=True,
-                return_tensors="pt",
-            )
-
-            text_input_ids = text_inputs.input_ids
-            untruncated_ids = tokenizer(
-                prompt, padding="longest", return_tensors="pt"
-            ).input_ids
-
-            if untruncated_ids.shape[-1] >= text_input_ids.shape[
-                -1
-            ] and not torch.equal(text_input_ids, untruncated_ids):
-                removed_text = tokenizer.batch_decode(
-                    untruncated_ids[:, tokenizer.model_max_length - 1 : -1]
-                )
-                print(
-                    "The following part of your input was truncated because CLIP can only handle sequences up to"
-                    f" {tokenizer.model_max_length} tokens: {removed_text}"
-                )
-
-            text_encoder_output = text_encoder("forward", (text_input_ids,))
-            prompt_embeds = torch.from_numpy(text_encoder_output[0])
-            pooled_prompt_embeds = torch.from_numpy(text_encoder_output[1])
-
-            prompt_embeds_list.append(prompt_embeds)
-
-        prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
-
-        # get unconditional embeddings for classifier free guidance
-        zero_out_negative_prompt = (
-            negative_prompt is None
-            and self.config.force_zeros_for_empty_prompt
-        )
-        if (
-            do_classifier_free_guidance
-            and negative_prompt_embeds is None
-            and zero_out_negative_prompt
-        ):
-            negative_prompt_embeds = torch.zeros_like(prompt_embeds)
-            negative_pooled_prompt_embeds = torch.zeros_like(
-                pooled_prompt_embeds
-            )
-        elif do_classifier_free_guidance and negative_prompt_embeds is None:
-            negative_prompt = negative_prompt or ""
-            negative_prompt_2 = negative_prompt
-
-            uncond_tokens: List[str]
-            if prompt is not None and type(prompt) is not type(
-                negative_prompt
-            ):
-                raise TypeError(
-                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    f" {type(prompt)}."
-                )
-            elif isinstance(negative_prompt, str):
-                uncond_tokens = [negative_prompt, negative_prompt_2]
-            elif batch_size != len(negative_prompt):
-                raise ValueError(
-                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
-                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
-                    " the batch size of `prompt`."
-                )
-            else:
-                uncond_tokens = [negative_prompt, negative_prompt_2]
-
-            negative_prompt_embeds_list = []
-            for negative_prompt, tokenizer, text_encoder in zip(
-                uncond_tokens, tokenizers, text_encoders
-            ):
-                max_length = prompt_embeds.shape[1]
-                uncond_input = tokenizer(
-                    negative_prompt,
-                    padding="max_length",
-                    max_length=max_length,
-                    truncation=True,
-                    return_tensors="pt",
-                )
-                text_encoder_output = text_encoder(
-                    "forward", (uncond_input.input_ids,)
-                )
-                negative_prompt_embeds = torch.from_numpy(
-                    text_encoder_output[0]
-                )
-                negative_pooled_prompt_embeds = torch.from_numpy(
-                    text_encoder_output[1]
-                )
-
-                negative_prompt_embeds_list.append(negative_prompt_embeds)
-
-            negative_prompt_embeds = torch.concat(
-                negative_prompt_embeds_list, dim=-1
-            )
-
-        if self.ondemand:
-            self.unload_clip_sdxl()
-            gc.collect()
-
-        # TODO: Look into dtype for text_encoder_2!
-        prompt_embeds = prompt_embeds.to(dtype=torch.float16)
-        bs_embed, seq_len, _ = prompt_embeds.shape
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(
-            bs_embed * num_images_per_prompt, seq_len, -1
-        )
-
-        # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
-        seq_len = negative_prompt_embeds.shape[1]
-        negative_prompt_embeds = negative_prompt_embeds.to(dtype=torch.float32)
-        negative_prompt_embeds = negative_prompt_embeds.repeat(
-            1, num_images_per_prompt, 1
-        )
-        negative_prompt_embeds = negative_prompt_embeds.view(
-            batch_size * num_images_per_prompt, seq_len, -1
-        )
-
-        pooled_prompt_embeds = pooled_prompt_embeds.repeat(
-            1, num_images_per_prompt
-        ).view(bs_embed * num_images_per_prompt, -1)
-        negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(
-            1, num_images_per_prompt
-        ).view(bs_embed * num_images_per_prompt, -1)
-
-        return (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        )
 
     def encode_prompts(self, prompts, neg_prompts, max_length):
         # Tokenize text and get embeddings
@@ -399,7 +177,6 @@ class StableDiffusionPipeline:
         clip_inf_time = (time.time() - clip_inf_start) * 1000
         if self.ondemand:
             self.unload_clip()
-            gc.collect()
         self.log += f"\nClip Inference time (ms) = {clip_inf_time:.3f}"
 
         return text_embeddings
@@ -512,8 +289,6 @@ class StableDiffusionPipeline:
         if self.ondemand:
             self.unload_unet()
             self.unload_unet_512()
-            gc.collect()
-
         avg_step_time = step_time_sum / len(total_timesteps)
         self.log += f"\nAverage step time: {avg_step_time}ms/it"
 
@@ -522,96 +297,6 @@ class StableDiffusionPipeline:
         all_latents = torch.cat(latent_history, dim=0)
         return all_latents
 
-    def produce_img_latents_sdxl(
-        self,
-        latents,
-        total_timesteps,
-        add_text_embeds,
-        add_time_ids,
-        prompt_embeds,
-        cpu_scheduling,
-        guidance_scale,
-        dtype,
-        mask=None,
-        masked_image_latents=None,
-        return_all_latents=False,
-    ):
-        # return None
-        self.status = SD_STATE_IDLE
-        step_time_sum = 0
-        extra_step_kwargs = {"generator": None}
-        self.load_unet()
-        for i, t in tqdm(enumerate(total_timesteps)):
-            step_start_time = time.time()
-            timestep = torch.tensor([t]).to(dtype).detach().numpy()
-            # expand the latents if we are doing classifier free guidance
-            if isinstance(latents, np.ndarray):
-                latents = torch.tensor(latents)
-            latent_model_input = torch.cat([latents] * 2)
-
-            latent_model_input = self.scheduler.scale_model_input(
-                latent_model_input, t
-            )
-            if mask is not None and masked_image_latents is not None:
-                latent_model_input = torch.cat(
-                    [
-                        torch.from_numpy(np.asarray(latent_model_input)),
-                        mask,
-                        masked_image_latents,
-                    ],
-                    dim=1,
-                ).to(dtype)
-
-            noise_pred = self.unet(
-                "forward",
-                (
-                    latent_model_input,
-                    timestep,
-                    prompt_embeds,
-                    add_text_embeds,
-                    add_time_ids,
-                    guidance_scale,
-                ),
-                send_to_host=True,
-            )
-            if not isinstance(latents, torch.Tensor):
-                latents = torch.from_numpy(latents).to("cpu")
-            noise_pred = torch.from_numpy(noise_pred).to("cpu")
-
-            latents = self.scheduler.step(
-                noise_pred, t, latents, **extra_step_kwargs, return_dict=False
-            )[0]
-            latents = latents.detach().numpy()
-            noise_pred = noise_pred.detach().numpy()
-
-            step_time = (time.time() - step_start_time) * 1000
-            step_time_sum += step_time
-
-            if self.status == SD_STATE_CANCEL:
-                break
-        if self.ondemand:
-            self.unload_unet()
-            gc.collect()
-
-        avg_step_time = step_time_sum / len(total_timesteps)
-        self.log += f"\nAverage step time: {avg_step_time}ms/it"
-
-        return latents
-
-    def decode_latents_sdxl(self, latents, is_fp32_vae):
-        # latents are in unet dtype here so switch if we want to use fp32
-        if is_fp32_vae:
-            print("Casting latents to float32 for VAE")
-            latents = latents.to(torch.float32)
-        images = self.vae("forward", (latents,))
-        images = (torch.from_numpy(images) / 2 + 0.5).clamp(0, 1)
-        images = images.cpu().permute(0, 2, 3, 1).float().numpy()
-
-        images = (images * 255).round().astype("uint8")
-        pil_images = [Image.fromarray(image[:, :, :3]) for image in images]
-
-        return pil_images
-
     @classmethod
     def from_pretrained(
         cls,
@@ -625,10 +310,6 @@ class StableDiffusionPipeline:
             DPMSolverMultistepScheduler,
             SharkEulerDiscreteScheduler,
             DEISMultistepScheduler,
-            DDPMScheduler,
-            DPMSolverSinglestepScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            HeunDiscreteScheduler,
         ],
         import_mlir: bool,
         model_id: str,
@@ -644,8 +325,7 @@ class StableDiffusionPipeline:
         ondemand: bool,
         low_cpu_mem_usage: bool = False,
         debug: bool = False,
-        stencils: list[str] = [],
-        # stencil_images: list[Image] = []
+        use_stencil: str = None,
         use_lora: str = "",
         ddpm_scheduler: DDPMScheduler = None,
         use_quantize=None,
@@ -662,7 +342,6 @@ class StableDiffusionPipeline:
             "OutpaintPipeline",
         ]
         is_upscaler = cls.__name__ in ["UpscalerPipeline"]
-        is_sdxl = cls.__name__ in ["Text2ImageSDXLPipeline"]
 
         sd_model = SharkifyStableDiffusionModel(
             model_id,
@@ -679,8 +358,7 @@ class StableDiffusionPipeline:
             debug=debug,
             is_inpaint=is_inpaint,
             is_upscaler=is_upscaler,
-            is_sdxl=is_sdxl,
-            stencils=stencils,
+            use_stencil=use_stencil,
             use_lora=use_lora,
             use_quantize=use_quantize,
         )
@@ -695,21 +373,6 @@ class StableDiffusionPipeline:
                 ondemand,
             )
 
-        if cls.__name__ == "StencilPipeline":
-            return cls(
-                scheduler, sd_model, import_mlir, use_lora, ondemand, stencils
-            )
-        if cls.__name__ == "Text2ImageSDXLPipeline":
-            is_fp32_vae = True if "16" not in custom_vae else False
-            return cls(
-                scheduler,
-                sd_model,
-                import_mlir,
-                use_lora,
-                ondemand,
-                is_fp32_vae,
-            )
-
         return cls(scheduler, sd_model, import_mlir, use_lora, ondemand)
 
     # #####################################################
@@ -731,21 +394,16 @@ class StableDiffusionPipeline:
             prompt (`str` or `list(int)`):
                 prompt to be encoded
             negative_prompt (`str` or `List[str]`):
-                The prompt or prompts not to guide the image generation.
-                Ignored when not using guidance
-                (i.e., ignored if `guidance_scale` is less than `1`).
+                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
+                if `guidance_scale` is less than `1`).
             model_max_length (int):
-                SHARK: pass the max length instead of relying on
-                pipe.tokenizer.model_max_length
+                SHARK: pass the max length instead of relying on pipe.tokenizer.model_max_length
             do_classifier_free_guidance (`bool`):
                 whether to use classifier free guidance or not,
-                SHARK: must be set to True as we always expect neg embeddings
-                (defaulted to True)
+                SHARK: must be set to True as we always expect neg embeddings (defaulted to True)
             max_embeddings_multiples (`int`, *optional*, defaults to `3`):
-                The max multiple length of prompt embeddings compared to the
-                max output length of text encoder.
-                SHARK: max_embeddings_multiples>1 produce a tensor shape error
-                (defaulted to 1)
+                The max multiple length of prompt embeddings compared to the max output length of text encoder.
+                SHARK: max_embeddings_multiples>1 produce a tensor shape error (defaulted to 1)
             num_images_per_prompt (`int`):
                 number of images that should be generated per prompt
                 SHARK: num_images_per_prompt is not used (defaulted to 1)
@@ -764,11 +422,9 @@ class StableDiffusionPipeline:
             negative_prompt = [negative_prompt] * batch_size
         if batch_size != len(negative_prompt):
             raise ValueError(
-                f"`negative_prompt`: "
-                f"{negative_prompt} has batch size {len(negative_prompt)}, "
-                f"but `prompt`: {prompt} has batch size {batch_size}. "
-                f"Please make sure that passed `negative_prompt` matches "
-                "the batch size of `prompt`."
+                f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                " the batch size of `prompt`."
             )
 
         text_embeddings, uncond_embeddings = get_weighted_text_embeddings(
@@ -781,36 +437,14 @@ class StableDiffusionPipeline:
         )
         # SHARK: we are not using num_images_per_prompt
         # bs_embed, seq_len, _ = text_embeddings.shape
-        # text_embeddings = text_embeddings.repeat(
-        #     1,
-        #     num_images_per_prompt,
-        #     1
-        # )
-        # text_embeddings = (
-        #     text_embeddings.view(
-        #         bs_embed * num_images_per_prompt,
-        #         seq_len,
-        #         -1
-        #     )
-        # )
+        # text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)
+        # text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
 
         if do_classifier_free_guidance:
             # SHARK: we are not using num_images_per_prompt
             # bs_embed, seq_len, _ = uncond_embeddings.shape
-            # uncond_embeddings = (
-            #     uncond_embeddings.repeat(
-            #         1,
-            #         num_images_per_prompt,
-            #         1
-            #     )
-            # )
-            # uncond_embeddings = (
-            #     uncond_embeddings.view(
-            #         bs_embed * num_images_per_prompt,
-            #         seq_len,
-            #         -1
-            #     )
-            # )
+            # uncond_embeddings = uncond_embeddings.repeat(1, num_images_per_prompt, 1)
+            # uncond_embeddings = uncond_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
             text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
 
         if text_embeddings.shape[1] > model_max_length:
@@ -822,10 +456,9 @@ class StableDiffusionPipeline:
         clip_inf_time = (time.time() - clip_inf_start) * 1000
         if self.ondemand:
             self.unload_clip()
-            gc.collect()
         self.log += f"\nClip Inference time (ms) = {clip_inf_time:.3f}"
 
-        return text_embeddings.numpy().astype(np.float16)
+        return text_embeddings.numpy()
 
 
 from typing import List, Optional, Union
@@ -853,8 +486,7 @@ re_attention = re.compile(
 
 def parse_prompt_attention(text):
     """
-    Parses a string with attention tokens and returns a list of pairs:
-        text and its associated weight.
+    Parses a string with attention tokens and returns a list of pairs: text and its associated weight.
     Accepted tokens are:
       (abc) - increases attention to abc by a multiplier of 1.1
       (abc:3.12) - increases attention to abc by a multiplier of 3.12
@@ -1081,12 +713,6 @@ def get_unweighted_text_embeddings(
     return text_embeddings
 
 
-# This function deals with NoneType values occuring in tokens after padding
-# It switches out None with 49407 as truncating None values causes matrix dimension errors,
-def filter_nonetype_tokens(tokens: List[List]):
-    return [[49407 if token is None else token for token in tokens[0]]]
-
-
 def get_weighted_text_embeddings(
     pipe: StableDiffusionPipeline,
     prompt: Union[str, List[str]],
@@ -1178,10 +804,6 @@ def get_weighted_text_embeddings(
         no_boseos_middle=no_boseos_middle,
         chunk_length=pipe.model_max_length,
     )
-
-    # FIXME: This is a hacky fix caused by tokenizer padding with None values
-    prompt_tokens = filter_nonetype_tokens(prompt_tokens)
-
     # prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.long, device=pipe.device)
     prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.long, device="cpu")
     if uncond_prompt is not None:
@@ -1194,10 +816,6 @@ def get_weighted_text_embeddings(
             no_boseos_middle=no_boseos_middle,
             chunk_length=pipe.model_max_length,
         )
-
-        # FIXME: This is a hacky fix caused by tokenizer padding with None values
-        uncond_tokens = filter_nonetype_tokens(uncond_tokens)
-
         # uncond_tokens = torch.tensor(uncond_tokens, dtype=torch.long, device=pipe.device)
         uncond_tokens = torch.tensor(
             uncond_tokens, dtype=torch.long, device="cpu"

apps/stable_diffusion/src/schedulers/__init__.py

@@ -1,7 +1,4 @@
+from apps.stable_diffusion.src.schedulers.sd_schedulers import get_schedulers
 from apps.stable_diffusion.src.schedulers.shark_eulerdiscrete import (
     SharkEulerDiscreteScheduler,
 )
-from apps.stable_diffusion.src.schedulers.shark_eulerancestraldiscrete import (
-    SharkEulerAncestralDiscreteScheduler,
-)
-from apps.stable_diffusion.src.schedulers.sd_schedulers import get_schedulers

apps/stable_diffusion/src/schedulers/sd_schedulers.py

@@ -1,5 +1,4 @@
 from diffusers import (
-    LCMScheduler,
     LMSDiscreteScheduler,
     PNDMScheduler,
     DDPMScheduler,
@@ -9,28 +8,13 @@ from diffusers import (
     EulerDiscreteScheduler,
     EulerAncestralDiscreteScheduler,
     DEISMultistepScheduler,
-    DPMSolverSinglestepScheduler,
-    KDPM2AncestralDiscreteScheduler,
-    HeunDiscreteScheduler,
 )
 from apps.stable_diffusion.src.schedulers.shark_eulerdiscrete import (
     SharkEulerDiscreteScheduler,
 )
-from apps.stable_diffusion.src.schedulers.shark_eulerancestraldiscrete import (
-    SharkEulerAncestralDiscreteScheduler,
-)
 
 
 def get_schedulers(model_id):
-    # TODO: Robust scheduler setup on pipeline creation -- if we don't
-    # set batch_size here, the SHARK schedulers will
-    # compile with batch size = 1 regardless of whether the model
-    # outputs latents of a larger batch size, e.g. SDXL.
-    # However, obviously, searching for whether the base model ID
-    # contains "xl" is not very robust.
-
-    batch_size = 2 if "xl" in model_id.lower() else 1
-
     schedulers = dict()
     schedulers["PNDM"] = PNDMScheduler.from_pretrained(
         model_id,
@@ -52,34 +36,11 @@ def get_schedulers(model_id):
         model_id,
         subfolder="scheduler",
     )
-    schedulers["LCMScheduler"] = LCMScheduler.from_pretrained(
-        model_id,
-        subfolder="scheduler",
-    )
     schedulers[
         "DPMSolverMultistep"
-    ] = DPMSolverMultistepScheduler.from_pretrained(
-        model_id, subfolder="scheduler", algorithm_type="dpmsolver"
-    )
-    schedulers[
-        "DPMSolverMultistep++"
-    ] = DPMSolverMultistepScheduler.from_pretrained(
-        model_id, subfolder="scheduler", algorithm_type="dpmsolver++"
-    )
-    schedulers[
-        "DPMSolverMultistepKarras"
     ] = DPMSolverMultistepScheduler.from_pretrained(
         model_id,
         subfolder="scheduler",
-        use_karras_sigmas=True,
-    )
-    schedulers[
-        "DPMSolverMultistepKarras++"
-    ] = DPMSolverMultistepScheduler.from_pretrained(
-        model_id,
-        subfolder="scheduler",
-        algorithm_type="dpmsolver++",
-        use_karras_sigmas=True,
     )
     schedulers["EulerDiscrete"] = EulerDiscreteScheduler.from_pretrained(
         model_id,
@@ -101,28 +62,5 @@ def get_schedulers(model_id):
         model_id,
         subfolder="scheduler",
     )
-    schedulers[
-        "SharkEulerAncestralDiscrete"
-    ] = SharkEulerAncestralDiscreteScheduler.from_pretrained(
-        model_id,
-        subfolder="scheduler",
-    )
-    schedulers[
-        "DPMSolverSinglestep"
-    ] = DPMSolverSinglestepScheduler.from_pretrained(
-        model_id,
-        subfolder="scheduler",
-    )
-    schedulers[
-        "KDPM2AncestralDiscrete"
-    ] = KDPM2AncestralDiscreteScheduler.from_pretrained(
-        model_id,
-        subfolder="scheduler",
-    )
-    schedulers["HeunDiscrete"] = HeunDiscreteScheduler.from_pretrained(
-        model_id,
-        subfolder="scheduler",
-    )
-    schedulers["SharkEulerDiscrete"].compile(batch_size)
-    schedulers["SharkEulerAncestralDiscrete"].compile(batch_size)
+    schedulers["SharkEulerDiscrete"].compile()
     return schedulers

apps/stable_diffusion/src/schedulers/shark_eulerancestraldiscrete.py

@@ -1,251 +0,0 @@
-import sys
-import numpy as np
-from typing import List, Optional, Tuple, Union
-from diffusers import (
-    EulerAncestralDiscreteScheduler,
-)
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.configuration_utils import register_to_config
-from apps.stable_diffusion.src.utils import (
-    compile_through_fx,
-    get_shark_model,
-    args,
-)
-import torch
-
-
-class SharkEulerAncestralDiscreteScheduler(EulerAncestralDiscreteScheduler):
-    @register_to_config
-    def __init__(
-        self,
-        num_train_timesteps: int = 1000,
-        beta_start: float = 0.0001,
-        beta_end: float = 0.02,
-        beta_schedule: str = "linear",
-        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
-        prediction_type: str = "epsilon",
-        timestep_spacing: str = "linspace",
-        steps_offset: int = 0,
-    ):
-        super().__init__(
-            num_train_timesteps,
-            beta_start,
-            beta_end,
-            beta_schedule,
-            trained_betas,
-            prediction_type,
-            timestep_spacing,
-            steps_offset,
-        )
-        # TODO: make it dynamic so we dont have to worry about batch size
-        self.batch_size = None
-        self.init_input_shape = None
-
-    def compile(self, batch_size=1):
-        SCHEDULER_BUCKET = "gs://shark_tank/stable_diffusion/schedulers"
-        device = args.device.split(":", 1)[0].strip()
-        self.batch_size = batch_size
-
-        model_input = {
-            "eulera": {
-                "output": torch.randn(
-                    batch_size, 4, args.height // 8, args.width // 8
-                ),
-                "latent": torch.randn(
-                    batch_size, 4, args.height // 8, args.width // 8
-                ),
-                "sigma": torch.tensor(1).to(torch.float32),
-                "sigma_from": torch.tensor(1).to(torch.float32),
-                "sigma_to": torch.tensor(1).to(torch.float32),
-                "noise": torch.randn(
-                    batch_size, 4, args.height // 8, args.width // 8
-                ),
-            },
-        }
-
-        example_latent = model_input["eulera"]["latent"]
-        example_output = model_input["eulera"]["output"]
-        example_noise = model_input["eulera"]["noise"]
-        if args.precision == "fp16":
-            example_latent = example_latent.half()
-            example_output = example_output.half()
-            example_noise = example_noise.half()
-        example_sigma = model_input["eulera"]["sigma"]
-        example_sigma_from = model_input["eulera"]["sigma_from"]
-        example_sigma_to = model_input["eulera"]["sigma_to"]
-
-        class ScalingModel(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(self, latent, sigma):
-                return latent / ((sigma**2 + 1) ** 0.5)
-
-        class SchedulerStepEpsilonModel(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(
-                self, noise_pred, latent, sigma, sigma_from, sigma_to, noise
-            ):
-                sigma_up = (
-                    sigma_to**2
-                    * (sigma_from**2 - sigma_to**2)
-                    / sigma_from**2
-                ) ** 0.5
-                sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
-                dt = sigma_down - sigma
-                pred_original_sample = latent - sigma * noise_pred
-                derivative = (latent - pred_original_sample) / sigma
-                prev_sample = latent + derivative * dt
-                return prev_sample + noise * sigma_up
-
-        class SchedulerStepVPredictionModel(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(
-                self, noise_pred, sigma, sigma_from, sigma_to, latent, noise
-            ):
-                sigma_up = (
-                    sigma_to**2
-                    * (sigma_from**2 - sigma_to**2)
-                    / sigma_from**2
-                ) ** 0.5
-                sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
-                dt = sigma_down - sigma
-                pred_original_sample = noise_pred * (
-                    -sigma / (sigma**2 + 1) ** 0.5
-                ) + (latent / (sigma**2 + 1))
-                derivative = (latent - pred_original_sample) / sigma
-                prev_sample = latent + derivative * dt
-                return prev_sample + noise * sigma_up
-
-        iree_flags = []
-        if len(args.iree_vulkan_target_triple) > 0:
-            iree_flags.append(
-                f"-iree-vulkan-target-triple={args.iree_vulkan_target_triple}"
-            )
-
-        def _import(self):
-            scaling_model = ScalingModel()
-            self.scaling_model, _ = compile_through_fx(
-                model=scaling_model,
-                inputs=(example_latent, example_sigma),
-                extended_model_name=f"euler_a_scale_model_input_{self.batch_size}_{args.height}_{args.width}_{device}_"
-                + args.precision,
-                extra_args=iree_flags,
-            )
-
-            pred_type_model_dict = {
-                "epsilon": SchedulerStepEpsilonModel(),
-                "v_prediction": SchedulerStepVPredictionModel(),
-            }
-            step_model = pred_type_model_dict[self.config.prediction_type]
-            self.step_model, _ = compile_through_fx(
-                step_model,
-                (
-                    example_output,
-                    example_latent,
-                    example_sigma,
-                    example_sigma_from,
-                    example_sigma_to,
-                    example_noise,
-                ),
-                extended_model_name=f"euler_a_step_{self.config.prediction_type}_{self.batch_size}_{args.height}_{args.width}_{device}_"
-                + args.precision,
-                extra_args=iree_flags,
-            )
-
-        if args.import_mlir:
-            _import(self)
-
-        else:
-            try:
-                self.scaling_model = get_shark_model(
-                    SCHEDULER_BUCKET,
-                    "euler_a_scale_model_input_" + args.precision,
-                    iree_flags,
-                )
-                self.step_model = get_shark_model(
-                    SCHEDULER_BUCKET,
-                    "euler_a_step_"
-                    + self.config.prediction_type
-                    + args.precision,
-                    iree_flags,
-                )
-            except:
-                print(
-                    "failed to download model, falling back and using import_mlir"
-                )
-                args.import_mlir = True
-                _import(self)
-
-    def scale_model_input(self, sample, timestep):
-        if self.step_index is None:
-            self._init_step_index(timestep)
-        sigma = self.sigmas[self.step_index]
-        return self.scaling_model(
-            "forward",
-            (
-                sample,
-                sigma,
-            ),
-            send_to_host=False,
-        )
-
-    def step(
-        self,
-        noise_pred,
-        timestep,
-        latent,
-        generator: Optional[torch.Generator] = None,
-        return_dict: Optional[bool] = False,
-    ):
-        step_inputs = []
-
-        if self.step_index is None:
-            self._init_step_index(timestep)
-
-        sigma = self.sigmas[self.step_index]
-
-        sigma_from = self.sigmas[self.step_index]
-        sigma_to = self.sigmas[self.step_index + 1]
-        noise = randn_tensor(
-            torch.Size(noise_pred.shape),
-            dtype=torch.float16,
-            device="cpu",
-            generator=generator,
-        )
-        step_inputs = [
-            noise_pred,
-            latent,
-            sigma,
-            sigma_from,
-            sigma_to,
-            noise,
-        ]
-        # TODO: deal with dynamic inputs in turbine flow.
-        # update step index since we're done with the variable and will return with compiled module output.
-        self._step_index += 1
-
-        if noise_pred.shape[0] < self.batch_size:
-            for i in [0, 1, 5]:
-                try:
-                    step_inputs[i] = torch.tensor(step_inputs[i])
-                except:
-                    step_inputs[i] = torch.tensor(step_inputs[i].to_host())
-                step_inputs[i] = torch.cat(
-                    (step_inputs[i], step_inputs[i]), axis=0
-                )
-            return self.step_model(
-                "forward",
-                tuple(step_inputs),
-                send_to_host=True,
-            )
-
-        return self.step_model(
-            "forward",
-            tuple(step_inputs),
-            send_to_host=False,
-        )

apps/stable_diffusion/src/schedulers/shark_eulerdiscrete.py

@@ -2,9 +2,12 @@ import sys
 import numpy as np
 from typing import List, Optional, Tuple, Union
 from diffusers import (
+    LMSDiscreteScheduler,
+    PNDMScheduler,
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
     EulerDiscreteScheduler,
 )
-from diffusers.utils.torch_utils import randn_tensor
 from diffusers.configuration_utils import register_to_config
 from apps.stable_diffusion.src.utils import (
     compile_through_fx,
@@ -24,13 +27,6 @@ class SharkEulerDiscreteScheduler(EulerDiscreteScheduler):
         beta_schedule: str = "linear",
         trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
         prediction_type: str = "epsilon",
-        interpolation_type: str = "linear",
-        use_karras_sigmas: bool = False,
-        sigma_min: Optional[float] = None,
-        sigma_max: Optional[float] = None,
-        timestep_spacing: str = "linspace",
-        timestep_type: str = "discrete",
-        steps_offset: int = 0,
     ):
         super().__init__(
             num_train_timesteps,
@@ -39,29 +35,20 @@ class SharkEulerDiscreteScheduler(EulerDiscreteScheduler):
             beta_schedule,
             trained_betas,
             prediction_type,
-            interpolation_type,
-            use_karras_sigmas,
-            sigma_min,
-            sigma_max,
-            timestep_spacing,
-            timestep_type,
-            steps_offset,
         )
-        # TODO: make it dynamic so we dont have to worry about batch size
-        self.batch_size = 1
 
-    def compile(self, batch_size=1):
+    def compile(self):
         SCHEDULER_BUCKET = "gs://shark_tank/stable_diffusion/schedulers"
+        BATCH_SIZE = args.batch_size
         device = args.device.split(":", 1)[0].strip()
-        self.batch_size = batch_size
 
         model_input = {
             "euler": {
                 "latent": torch.randn(
-                    batch_size, 4, args.height // 8, args.width // 8
+                    BATCH_SIZE, 4, args.height // 8, args.width // 8
                 ),
                 "output": torch.randn(
-                    batch_size, 4, args.height // 8, args.width // 8
+                    BATCH_SIZE, 4, args.height // 8, args.width // 8
                 ),
                 "sigma": torch.tensor(1).to(torch.float32),
                 "dt": torch.tensor(1).to(torch.float32),
@@ -83,32 +70,12 @@ class SharkEulerDiscreteScheduler(EulerDiscreteScheduler):
             def forward(self, latent, sigma):
                 return latent / ((sigma**2 + 1) ** 0.5)
 
-        class SchedulerStepEpsilonModel(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(self, noise_pred, sigma_hat, latent, dt):
-                pred_original_sample = latent - sigma_hat * noise_pred
-                derivative = (latent - pred_original_sample) / sigma_hat
-                return latent + derivative * dt
-
-        class SchedulerStepSampleModel(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(self, noise_pred, sigma_hat, latent, dt):
-                pred_original_sample = noise_pred
-                derivative = (latent - pred_original_sample) / sigma_hat
-                return latent + derivative * dt
-
-        class SchedulerStepVPredictionModel(torch.nn.Module):
+        class SchedulerStepModel(torch.nn.Module):
             def __init__(self):
                 super().__init__()
 
             def forward(self, noise_pred, sigma, latent, dt):
-                pred_original_sample = noise_pred * (
-                    -sigma / (sigma**2 + 1) ** 0.5
-                ) + (latent / (sigma**2 + 1))
+                pred_original_sample = latent - sigma * noise_pred
                 derivative = (latent - pred_original_sample) / sigma
                 return latent + derivative * dt
 
@@ -117,28 +84,25 @@ class SharkEulerDiscreteScheduler(EulerDiscreteScheduler):
             iree_flags.append(
                 f"-iree-vulkan-target-triple={args.iree_vulkan_target_triple}"
             )
+        # Disable bindings fusion to work with moltenVK.
+        if sys.platform == "darwin":
+            iree_flags.append("-iree-stream-fuse-binding=false")
 
         def _import(self):
             scaling_model = ScalingModel()
             self.scaling_model, _ = compile_through_fx(
                 model=scaling_model,
                 inputs=(example_latent, example_sigma),
-                extended_model_name=f"euler_scale_model_input_{self.batch_size}_{args.height}_{args.width}_{device}_"
+                extended_model_name=f"euler_scale_model_input_{BATCH_SIZE}_{args.height}_{args.width}_{device}_"
                 + args.precision,
                 extra_args=iree_flags,
             )
 
-            pred_type_model_dict = {
-                "epsilon": SchedulerStepEpsilonModel(),
-                "v_prediction": SchedulerStepVPredictionModel(),
-                "sample": SchedulerStepSampleModel(),
-                "original_sample": SchedulerStepSampleModel(),
-            }
-            step_model = pred_type_model_dict[self.config.prediction_type]
+            step_model = SchedulerStepModel()
             self.step_model, _ = compile_through_fx(
                 step_model,
                 (example_output, example_sigma, example_latent, example_dt),
-                extended_model_name=f"euler_step_{self.config.prediction_type}_{self.batch_size}_{args.height}_{args.width}_{device}_"
+                extended_model_name=f"euler_step_{BATCH_SIZE}_{args.height}_{args.width}_{device}_"
                 + args.precision,
                 extra_args=iree_flags,
             )
@@ -148,11 +112,6 @@ class SharkEulerDiscreteScheduler(EulerDiscreteScheduler):
 
         else:
             try:
-                step_model_type = (
-                    "sample"
-                    if "sample" in self.config.prediction_type
-                    else self.config.prediction_type
-                )
                 self.scaling_model = get_shark_model(
                     SCHEDULER_BUCKET,
                     "euler_scale_model_input_" + args.precision,
@@ -160,7 +119,7 @@ class SharkEulerDiscreteScheduler(EulerDiscreteScheduler):
                 )
                 self.step_model = get_shark_model(
                     SCHEDULER_BUCKET,
-                    "euler_step_" + step_model_type + args.precision,
+                    "euler_step_" + args.precision,
                     iree_flags,
                 )
             except:
@@ -171,9 +130,8 @@ class SharkEulerDiscreteScheduler(EulerDiscreteScheduler):
                 _import(self)
 
     def scale_model_input(self, sample, timestep):
-        if self.step_index is None:
-            self._init_step_index(timestep)
-        sigma = self.sigmas[self.step_index]
+        step_index = (self.timesteps == timestep).nonzero().item()
+        sigma = self.sigmas[step_index]
         return self.scaling_model(
             "forward",
             (
@@ -183,61 +141,15 @@ class SharkEulerDiscreteScheduler(EulerDiscreteScheduler):
             send_to_host=False,
         )
 
-    def step(
-        self,
-        noise_pred,
-        timestep,
-        latent,
-        s_churn: float = 0.0,
-        s_tmin: float = 0.0,
-        s_tmax: float = float("inf"),
-        s_noise: float = 1.0,
-        generator: Optional[torch.Generator] = None,
-        return_dict: Optional[bool] = False,
-    ):
-        if self.step_index is None:
-            self._init_step_index(timestep)
-
-        sigma = self.sigmas[self.step_index]
-
-        gamma = (
-            min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1)
-            if s_tmin <= sigma <= s_tmax
-            else 0.0
-        )
-
-        sigma_hat = sigma * (gamma + 1)
-
-        noise_pred = (
-            torch.from_numpy(noise_pred)
-            if isinstance(noise_pred, np.ndarray)
-            else noise_pred
-        )
-
-        noise = randn_tensor(
-            torch.Size(noise_pred.shape),
-            dtype=torch.float16,
-            device="cpu",
-            generator=generator,
-        )
-
-        eps = noise * s_noise
-
-        if gamma > 0:
-            latent = latent + eps * (sigma_hat**2 - sigma**2) ** 0.5
-
-        if self.config.prediction_type == "v_prediction":
-            sigma_hat = sigma
-
-        dt = self.sigmas[self.step_index + 1] - sigma_hat
-
-        self._step_index += 1
-
+    def step(self, noise_pred, timestep, latent):
+        step_index = (self.timesteps == timestep).nonzero().item()
+        sigma = self.sigmas[step_index]
+        dt = self.sigmas[step_index + 1] - sigma
         return self.step_model(
             "forward",
             (
                 noise_pred,
-                sigma_hat,
+                sigma,
                 latent,
                 dt,
             ),

apps/stable_diffusion/src/utils/__init__.py

@@ -13,7 +13,6 @@ from apps.stable_diffusion.src.utils.sd_annotation import sd_model_annotation
 from apps.stable_diffusion.src.utils.stable_args import args
 from apps.stable_diffusion.src.utils.stencils.stencil_utils import (
     controlnet_hint_conversion,
-    controlnet_hint_reshaping,
     get_stencil_model_id,
 )
 from apps.stable_diffusion.src.utils.utils import (
@@ -29,8 +28,6 @@ from apps.stable_diffusion.src.utils.utils import (
     fetch_and_update_base_model_id,
     get_path_to_diffusers_checkpoint,
     sanitize_seed,
-    parse_seed_input,
-    batch_seeds,
     get_path_stem,
     get_extended_name,
     get_generated_imgs_path,
@@ -42,8 +39,3 @@ from apps.stable_diffusion.src.utils.utils import (
     resize_stencil,
     _compile_module,
 )
-from apps.stable_diffusion.src.utils.civitai import get_civitai_checkpoint
-from apps.stable_diffusion.src.utils.resamplers import (
-    resamplers,
-    resampler_list,
-)

apps/stable_diffusion/src/utils/civitai.py

@@ -1,42 +0,0 @@
-import re
-import requests
-from apps.stable_diffusion.src.utils.stable_args import args
-
-from pathlib import Path
-from tqdm import tqdm
-
-
-def get_civitai_checkpoint(url: str):
-    with requests.get(url, allow_redirects=True, stream=True) as response:
-        response.raise_for_status()
-
-        # civitai api returns the filename in the content disposition
-        base_filename = re.findall(
-            '"([^"]*)"', response.headers["Content-Disposition"]
-        )[0]
-        destination_path = (
-            Path.cwd() / (args.ckpt_dir or "models") / base_filename
-        )
-
-        # we don't have this model downloaded yet
-        if not destination_path.is_file():
-            print(
-                f"downloading civitai model from {url} to {destination_path}"
-            )
-
-            size = int(response.headers["content-length"], 0)
-            progress_bar = tqdm(total=size, unit="iB", unit_scale=True)
-
-            with open(destination_path, "wb") as f:
-                for chunk in response.iter_content(chunk_size=65536):
-                    f.write(chunk)
-                    progress_bar.update(len(chunk))
-
-            progress_bar.close()
-
-        # we already have this model downloaded
-        else:
-            print(f"civitai model already downloaded to {destination_path}")
-
-        response.close()
-        return destination_path.as_posix()

apps/stable_diffusion/src/utils/profiler.py

@@ -3,9 +3,7 @@ from apps.stable_diffusion.src.utils.stable_args import args
 
 # Helper function to profile the vulkan device.
 def start_profiling(file_path="foo.rdc", profiling_mode="queue"):
-    from shark.parser import shark_args
-
-    if shark_args.vulkan_debug_utils and "vulkan" in args.device:
+    if args.vulkan_debug_utils and "vulkan" in args.device:
         import iree
 
         print(f"Profiling and saving to {file_path}.")

apps/stable_diffusion/src/utils/resamplers.py

@@ -1,12 +0,0 @@
-import PIL.Image as Image
-
-resamplers = {
-    "Lanczos": Image.Resampling.LANCZOS,
-    "Nearest Neighbor": Image.Resampling.NEAREST,
-    "Bilinear": Image.Resampling.BILINEAR,
-    "Bicubic": Image.Resampling.BICUBIC,
-    "Hamming": Image.Resampling.HAMMING,
-    "Box": Image.Resampling.BOX,
-}
-
-resampler_list = resamplers.keys()

apps/stable_diffusion/src/utils/resources/base_model.json

@@ -8,15 +8,6 @@
             "dtype":"i64"
         }
     },
-    "sdxl_clip": {
-        "token" : {
-            "shape" : [
-                "1*batch_size",
-                "max_len"
-            ],
-            "dtype":"i64"
-        }
-    },
     "vae_encode": {
         "image" : {
             "shape" : [
@@ -188,95 +179,9 @@
                 "shape": [2],
                 "dtype": "i64"
             }
-        },
-        "stabilityai/sdxl-turbo": {
-            "latents": {
-                "shape": [
-                    "2*batch_size",
-                    4,
-                    "height",
-                    "width"
-                ],
-                "dtype": "f32"
-            },
-            "timesteps": {
-                "shape": [
-                    1
-                ],
-                "dtype": "f32"
-            },
-            "prompt_embeds": {
-                "shape": [
-                    "2*batch_size",
-                    "max_len",
-                    2048
-                ],
-                "dtype": "f32"
-            },
-            "text_embeds": {
-                "shape": [
-                    "2*batch_size",
-                    1280
-                ],
-                "dtype": "f32"
-            },
-            "time_ids": {
-                "shape": [
-                    "2*batch_size",
-                    6
-                ],
-                "dtype": "f32"
-            },
-            "guidance_scale": {
-                "shape": 1,
-                "dtype": "f32"
-            }
-        },
-        "stabilityai/stable-diffusion-xl-base-1.0": {
-            "latents": {
-                "shape": [
-                    "2*batch_size",
-                    4,
-                    "height",
-                    "width"
-                ],
-                "dtype": "f32"
-            },
-            "timesteps": {
-                "shape": [
-                    1
-                ],
-                "dtype": "f32"
-            },
-            "prompt_embeds": {
-                "shape": [
-                    "2*batch_size",
-                    "max_len",
-                    2048
-                ],
-                "dtype": "f32"
-            },
-            "text_embeds": {
-                "shape": [
-                    "2*batch_size",
-                    1280
-                ],
-                "dtype": "f32"
-            },
-            "time_ids": {
-                "shape": [
-                    "2*batch_size",
-                    6
-                ],
-                "dtype": "f32"
-            },
-            "guidance_scale": {
-                "shape": 1,
-                "dtype": "f32"
-            }
         }
     },
-    "stencil_adapter": {
+    "stencil_adaptor": {
         "latents": {
             "shape": [
                 "1*batch_size",
@@ -303,58 +208,6 @@
         "controlnet_hint": {
             "shape": [1, 3, "8*height", "8*width"],
             "dtype": "f32"
-        },
-        "acc1": {
-            "shape": [2, 320, "height", "width"],
-            "dtype": "f32"
-        },
-        "acc2": {
-            "shape": [2, 320, "height", "width"],
-            "dtype": "f32"
-        },
-        "acc3": {
-            "shape": [2, 320, "height", "width"],
-            "dtype": "f32"
-        },
-        "acc4": {
-            "shape": [2, 320, "height/2", "width/2"],
-            "dtype": "f32"
-        },
-        "acc5": {
-            "shape": [2, 640, "height/2", "width/2"],
-            "dtype": "f32"
-        },
-        "acc6": {
-            "shape": [2, 640, "height/2", "width/2"],
-            "dtype": "f32"
-        },
-        "acc7": {
-            "shape": [2, 640, "height/4", "width/4"],
-            "dtype": "f32"
-        },
-        "acc8": {
-            "shape": [2, 1280, "height/4", "width/4"],
-            "dtype": "f32"
-        },
-        "acc9": {
-            "shape": [2, 1280, "height/4", "width/4"],
-            "dtype": "f32"
-        },
-        "acc10": {
-            "shape": [2, 1280, "height/8", "width/8"],
-            "dtype": "f32"
-        },
-        "acc11": {
-            "shape": [2, 1280, "height/8", "width/8"],
-            "dtype": "f32"
-        },
-        "acc12": {
-            "shape": [2, 1280, "height/8", "width/8"],
-            "dtype": "f32"
-        },
-        "acc13": {
-            "shape": [2, 1280, "height/8", "width/8"],
-            "dtype": "f32"
         }
     },
     "stencil_unet": {
@@ -437,59 +290,7 @@
             "control13": {
                 "shape": [2, 1280, "height/8", "width/8"],
                 "dtype": "f32"
-            },
-            "scale1": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale2": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale3": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale4": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale5": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale6": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale7": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale8": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale9": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale10": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale11": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale12": {
-                "shape": 1,
-                "dtype": "f32"
-            },
-            "scale13": {
-                "shape": 1,
-                "dtype": "f32"
             }
         }
     }
-}
+}

apps/stable_diffusion/src/utils/resources/opt_flags.json

@@ -11,12 +11,12 @@
     "untuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       }
     }
@@ -28,7 +28,7 @@
         "specified_compilation_flags": {
           "cuda": [],
           "default_device": [
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
+            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
           ]
         }
       },
@@ -37,7 +37,7 @@
         "specified_compilation_flags": {
           "cuda": [],
           "default_device": [
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16},iree-linalg-ext-convert-conv2d-to-winograd))"
+            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16},iree-linalg-ext-convert-conv2d-to-winograd))"
           ]
         }
       }
@@ -45,12 +45,12 @@
     "untuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16},iree-linalg-ext-convert-conv2d-to-winograd))"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16},iree-linalg-ext-convert-conv2d-to-winograd))"
         ]
       }
     }
@@ -59,28 +59,24 @@
     "tuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))",
-          "--iree-opt-data-tiling=False"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))",
-          "--iree-opt-data-tiling=False"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       }
     },
     "untuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))",
-          "--iree-opt-data-tiling=False"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))",
-          "--iree-opt-data-tiling=False"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       }
     }

apps/stable_diffusion/src/utils/resources/prompts.json

@@ -1,12 +1,8 @@
 [["A high tech solarpunk utopia in the Amazon rainforest"],
-["Astrophotography, the shark nebula, nebula with a tiny shark-like cloud in the middle in the middle, hubble telescope, vivid colors"],
 ["A pikachu fine dining with a view to the Eiffel Tower"],
 ["A mecha robot in a favela in expressionist style"],
 ["an insect robot preparing a delicious meal"],
 ["A digital Illustration of the Babel tower, 4k, detailed, trending in artstation, fantasy vivid colors"],
 ["Cluttered house in the woods, anime, oil painting, high resolution, cottagecore, ghibli inspired, 4k"],
 ["A beautiful mansion beside a waterfall in the woods, by josef thoma, matte painting, trending on artstation HQ"],
-["portrait photo of a asia old warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes"],
-["A photo of a beach, sunset, calm, beautiful landscape, waves, water"],
-["(a large body of water with snowy mountains in the background), (fog, foggy, rolling fog), (clouds, cloudy, rolling clouds), dramatic sky and landscape, extraordinary landscape, (beautiful snow capped mountain background), (forest, dirt path)"],
-["a photo taken of the front of a super-car drifting on a road near mountains at high speeds with smokes coming off the tires, front angle, front point of view, trees in the mountains of the background, ((sharp focus))"]]
+["portrait photo of a asia old warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes"]]

apps/stable_diffusion/src/utils/sd_annotation.py

@@ -109,7 +109,7 @@ def load_lower_configs(base_model_id=None):
             spec = spec.split("-")[0]
 
     if args.annotation_model == "vae":
-        if not spec or spec in ["sm_80"]:
+        if not spec or spec in ["rdna3", "sm_80"]:
             config_name = (
                 f"{args.annotation_model}_{args.precision}_{device}.json"
             )
@@ -131,36 +131,15 @@ def load_lower_configs(base_model_id=None):
             "v1_4",
             "v1_5",
         ]:
-            config_name = (
-                f"{args.annotation_model}_"
-                f"{version}_"
-                f"{args.max_length}_"
-                f"{args.precision}_"
-                f"{device}_"
-                f"{spec}_"
-                f"{args.width}x{args.height}.json"
-            )
+            config_name = f"{args.annotation_model}_{version}_{args.max_length}_{args.precision}_{device}_{spec}_{args.width}x{args.height}.json"
         elif spec in ["rdna2"] and version in ["v2_1", "v2_1base", "v1_4"]:
-            config_name = (
-                f"{args.annotation_model}_"
-                f"{version}_"
-                f"{args.precision}_"
-                f"{device}_"
-                f"{spec}_"
-                f"{args.width}x{args.height}.json"
-            )
+            config_name = f"{args.annotation_model}_{version}_{args.precision}_{device}_{spec}_{args.width}x{args.height}.json"
         else:
-            config_name = (
-                f"{args.annotation_model}_"
-                f"{version}_"
-                f"{args.precision}_"
-                f"{device}_"
-                f"{spec}.json"
-            )
+            config_name = f"{args.annotation_model}_{version}_{args.precision}_{device}_{spec}.json"
 
+    full_gs_url = config_bucket + config_name
     lowering_config_dir = os.path.join(WORKDIR, "configs", config_name)
     print("Loading lowering config file from ", lowering_config_dir)
-    full_gs_url = config_bucket + config_name
     download_public_file(full_gs_url, lowering_config_dir, True)
     return lowering_config_dir
 
@@ -201,22 +180,9 @@ def dump_after_mlir(input_mlir, use_winograd):
 
     device, device_spec_args = get_device_args()
     if use_winograd:
-        preprocess_flag = (
-            "--iree-preprocessing-pass-pipeline=builtin.module"
-            "(func.func(iree-global-opt-detach-elementwise-from-named-ops,"
-            "iree-global-opt-convert-1x1-filter-conv2d-to-matmul,"
-            "iree-preprocessing-convert-conv2d-to-img2col,"
-            "iree-preprocessing-pad-linalg-ops{pad-size=32},"
-            "iree-linalg-ext-convert-conv2d-to-winograd))"
-        )
+        preprocess_flag = "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
     else:
-        preprocess_flag = (
-            "--iree-preprocessing-pass-pipeline=builtin.module"
-            "(func.func(iree-global-opt-detach-elementwise-from-named-ops,"
-            "iree-global-opt-convert-1x1-filter-conv2d-to-matmul,"
-            "iree-preprocessing-convert-conv2d-to-img2col,"
-            "iree-preprocessing-pad-linalg-ops{pad-size=32}))"
-        )
+        preprocess_flag = "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
 
     dump_module = ireec.compile_str(
         input_mlir,

apps/stable_diffusion/src/utils/stable_args.py

@@ -2,8 +2,6 @@ import argparse
 import os
 from pathlib import Path
 
-from apps.stable_diffusion.src.utils.resamplers import resampler_list
-
 
 def path_expand(s):
     return Path(s).expanduser().resolve()
@@ -21,56 +19,48 @@ p = argparse.ArgumentParser(
 )
 
 ##############################################################################
-# Stable Diffusion Params
+### Stable Diffusion Params
 ##############################################################################
 
 p.add_argument(
     "-a",
     "--app",
     default="txt2img",
-    help="Which app to use, one of: txt2img, img2img, outpaint, inpaint.",
+    help="which app to use, one of: txt2img, img2img, outpaint, inpaint",
 )
 p.add_argument(
     "-p",
     "--prompts",
     nargs="+",
-    default=[
-        "a photo taken of the front of a super-car drifting on a road near "
-        "mountains at high speeds with smokes coming off the tires, front "
-        "angle, front point of view, trees in the mountains of the "
-        "background, ((sharp focus))"
-    ],
-    help="Text of which images to be generated.",
+    default=["cyberpunk forest by Salvador Dali"],
+    help="text of which images to be generated.",
 )
 
 p.add_argument(
     "--negative_prompts",
     nargs="+",
-    default=[
-        "watermark, signature, logo, text, lowres, ((monochrome, grayscale)), "
-        "blurry, ugly, blur, oversaturated, cropped"
-    ],
-    help="Text you don't want to see in the generated image.",
+    default=["trees, green"],
+    help="text you don't want to see in the generated image.",
 )
 
 p.add_argument(
     "--img_path",
     type=str,
-    help="Path to the image input for img2img/inpainting.",
+    help="Path to the image input for img2img/inpainting",
 )
 
 p.add_argument(
     "--steps",
     type=int,
     default=50,
-    help="The number of steps to do the sampling.",
+    help="the no. of steps to do the sampling.",
 )
 
 p.add_argument(
     "--seed",
-    type=str,
+    type=int,
     default=-1,
-    help="The seed or list of seeds to use. -1 for a random one.",
+    help="the seed to use. -1 for a random one.",
 )
 
 p.add_argument(
@@ -78,143 +68,100 @@ p.add_argument(
     type=int,
     default=1,
     choices=range(1, 4),
-    help="The number of inferences to be made in a single `batch_count`.",
+    help="the number of inferences to be made in a single `batch_count`.",
 )
 
 p.add_argument(
     "--height",
     type=int,
     default=512,
-    choices=range(128, 1025, 8),
-    help="The height of the output image.",
+    choices=range(128, 769, 8),
+    help="the height of the output image.",
 )
 
 p.add_argument(
     "--width",
     type=int,
     default=512,
-    choices=range(128, 1025, 8),
-    help="The width of the output image.",
+    choices=range(128, 769, 8),
+    help="the width of the output image.",
 )
 
 p.add_argument(
     "--guidance_scale",
     type=float,
     default=7.5,
-    help="The value to be used for guidance scaling.",
+    help="the value to be used for guidance scaling.",
 )
 
 p.add_argument(
     "--noise_level",
     type=int,
     default=20,
-    help="The value to be used for noise level of upscaler.",
+    help="the value to be used for noise level of upscaler.",
 )
 
 p.add_argument(
     "--max_length",
     type=int,
     default=64,
-    help="Max length of the tokenizer output, options are 64 and 77.",
+    help="max length of the tokenizer output, options are 64 and 77.",
 )
 
 p.add_argument(
     "--max_embeddings_multiples",
     type=int,
     default=5,
-    help="The max multiple length of prompt embeddings compared to the max "
-    "output length of text encoder.",
+    help="The max multiple length of prompt embeddings compared to the max output length of text encoder.",
 )
 
 p.add_argument(
     "--strength",
     type=float,
     default=0.8,
-    help="The strength of change applied on the given input image for "
-    "img2img.",
-)
-
-p.add_argument(
-    "--use_hiresfix",
-    type=bool,
-    default=False,
-    help="Use Hires Fix to do higher resolution images, while trying to "
-    "avoid the issues that come with it. This is accomplished by first "
-    "generating an image using txt2img, then running it through img2img.",
-)
-
-p.add_argument(
-    "--hiresfix_height",
-    type=int,
-    default=768,
-    choices=range(128, 769, 8),
-    help="The height of the Hires Fix image.",
-)
-
-p.add_argument(
-    "--hiresfix_width",
-    type=int,
-    default=768,
-    choices=range(128, 769, 8),
-    help="The width of the Hires Fix image.",
-)
-
-p.add_argument(
-    "--hiresfix_strength",
-    type=float,
-    default=0.6,
-    help="The denoising strength to apply for the Hires Fix.",
-)
-
-p.add_argument(
-    "--resample_type",
-    type=str,
-    default="Nearest Neighbor",
-    choices=resampler_list,
-    help="The resample type to use when resizing an image before being run "
-    "through stable diffusion.",
+    help="the strength of change applied on the given input image for img2img",
 )
 
 ##############################################################################
-# Stable Diffusion Training Params
+### Stable Diffusion Training Params
 ##############################################################################
 
 p.add_argument(
     "--lora_save_dir",
     type=str,
     default="models/lora/",
-    help="Directory to save the lora fine tuned model.",
+    help="Directory to save the lora fine tuned model",
 )
 
 p.add_argument(
     "--training_images_dir",
     type=str,
     default="models/lora/training_images/",
-    help="Directory containing images that are an example of the prompt.",
+    help="Directory containing images that are an example of the prompt",
 )
 
 p.add_argument(
     "--training_steps",
     type=int,
     default=2000,
-    help="The number of steps to train.",
+    help="The no. of steps to train",
 )
 
 ##############################################################################
-# Inpainting and Outpainting Params
+### Inpainting and Outpainting Params
 ##############################################################################
 
 p.add_argument(
     "--mask_path",
     type=str,
-    help="Path to the mask image input for inpainting.",
+    help="Path to the mask image input for inpainting",
 )
 
 p.add_argument(
     "--inpaint_full_res",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="If inpaint only masked area or whole picture.",
+    help="If inpaint only masked area or whole picture",
 )
 
 p.add_argument(
@@ -222,7 +169,7 @@ p.add_argument(
     type=int,
     default=32,
     choices=range(0, 257, 4),
-    help="Number of pixels for only masked padding.",
+    help="Number of pixels for only masked padding",
 )
 
 p.add_argument(
@@ -230,7 +177,7 @@ p.add_argument(
     type=int,
     default=128,
     choices=range(8, 257, 8),
-    help="Number of expended pixels for one direction for outpainting.",
+    help="Number of expended pixels for one direction for outpainting",
 )
 
 p.add_argument(
@@ -238,94 +185,89 @@ p.add_argument(
     type=int,
     default=8,
     choices=range(0, 65),
-    help="Number of blur pixels for outpainting.",
+    help="Number of blur pixels for outpainting",
 )
 
 p.add_argument(
     "--left",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="If extend left for outpainting.",
+    help="If expend left for outpainting",
 )
 
 p.add_argument(
     "--right",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="If extend right for outpainting.",
+    help="If expend right for outpainting",
 )
 
 p.add_argument(
-    "--up",
     "--top",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="If extend top for outpainting.",
+    help="If expend top for outpainting",
 )
 
 p.add_argument(
-    "--down",
     "--bottom",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="If extend bottom for outpainting.",
+    help="If expend bottom for outpainting",
 )
 
 p.add_argument(
     "--noise_q",
     type=float,
     default=1.0,
-    help="Fall-off exponent for outpainting (lower=higher detail) "
-    "(min=0.0, max=4.0).",
+    help="Fall-off exponent for outpainting (lower=higher detail) (min=0.0, max=4.0)",
 )
 
 p.add_argument(
     "--color_variation",
     type=float,
     default=0.05,
-    help="Color variation for outpainting (min=0.0, max=1.0).",
+    help="Color variation for outpainting (min=0.0, max=1.0)",
 )
 
 ##############################################################################
-# Model Config and Usage Params
+### Model Config and Usage Params
 ##############################################################################
 
 p.add_argument(
-    "--device", type=str, default="vulkan", help="Device to run the model."
+    "--device", type=str, default="vulkan", help="device to run the model."
 )
 
 p.add_argument(
-    "--precision", type=str, default="fp16", help="Precision to run the model."
+    "--precision", type=str, default="fp16", help="precision to run the model."
 )
 
 p.add_argument(
     "--import_mlir",
-    default=True,
+    default=False,
     action=argparse.BooleanOptionalAction,
-    help="Imports the model from torch module to shark_module otherwise "
-    "downloads the model from shark_tank.",
+    help="imports the model from torch module to shark_module otherwise downloads the model from shark_tank.",
 )
 
 p.add_argument(
     "--load_vmfb",
     default=True,
     action=argparse.BooleanOptionalAction,
-    help="Attempts to load the model from a precompiled flat-buffer "
-    "and compiles + saves it if not found.",
+    help="attempts to load the model from a precompiled flatbuffer and compiles + saves it if not found.",
 )
 
 p.add_argument(
     "--save_vmfb",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="Saves the compiled flat-buffer to the local directory.",
+    help="saves the compiled flatbuffer to the local directory",
 )
 
 p.add_argument(
     "--use_tuned",
-    default=False,
+    default=True,
     action=argparse.BooleanOptionalAction,
-    help="Download and use the tuned version of the model if available.",
+    help="Download and use the tuned version of the model if available",
 )
 
 p.add_argument(
@@ -339,42 +281,28 @@ p.add_argument(
     "--scheduler",
     type=str,
     default="SharkEulerDiscrete",
-    help="Other supported schedulers are [DDIM, PNDM, LMSDiscrete, "
-    "DPMSolverMultistep, DPMSolverMultistep++, DPMSolverMultistepKarras, "
-    "DPMSolverMultistepKarras++, EulerDiscrete, EulerAncestralDiscrete, "
-    "DEISMultistep, KDPM2AncestralDiscrete, DPMSolverSinglestep, DDPM, "
-    "HeunDiscrete].",
+    help="other supported schedulers are [PNDM, DDIM, LMSDiscrete, EulerDiscrete, DPMSolverMultistep]",
 )
 
 p.add_argument(
     "--output_img_format",
     type=str,
     default="png",
-    help="Specify the format in which output image is save. "
-    "Supported options: jpg / png.",
+    help="specify the format in which output image is save. Supported options: jpg / png",
 )
 
 p.add_argument(
     "--output_dir",
     type=str,
     default=None,
-    help="Directory path to save the output images and json.",
+    help="Directory path to save the output images and json",
 )
 
 p.add_argument(
     "--batch_count",
     type=int,
     default=1,
-    help="Number of batches to be generated with random seeds in "
-    "single execution.",
-)
-
-p.add_argument(
-    "--repeatable_seeds",
-    default=False,
-    action=argparse.BooleanOptionalAction,
-    help="The seed of the first batch will be used as the rng seed to "
-    "generate the subsequent seeds for subsequent batches in that run.",
+    help="number of batch to be generated with random seeds in single execution",
 )
 
 p.add_argument(
@@ -388,8 +316,7 @@ p.add_argument(
     "--custom_vae",
     type=str,
     default="",
-    help="HuggingFace repo-id or path to SD model's checkpoint whose VAE "
-    "needs to be plugged in.",
+    help="HuggingFace repo-id or path to SD model's checkpoint whose Vae needs to be plugged in.",
 )
 
 p.add_argument(
@@ -403,304 +330,229 @@ p.add_argument(
     "--low_cpu_mem_usage",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="Use the accelerate package to reduce cpu memory consumption.",
+    help="Use the accelerate package to reduce cpu memory consumption",
 )
 
 p.add_argument(
     "--attention_slicing",
     type=str,
     default="none",
-    help="Amount of attention slicing to use (one of 'max', 'auto', 'none', "
-    "or an integer).",
+    help="Amount of attention slicing to use (one of 'max', 'auto', 'none', or an integer)",
 )
 
 p.add_argument(
     "--use_stencil",
-    choices=["canny", "openpose", "scribble", "zoedepth"],
+    choices=["canny", "openpose", "scribble"],
     help="Enable the stencil feature.",
 )
 
-p.add_argument(
-    "--control_mode",
-    choices=["Prompt", "Balanced", "Controlnet"],
-    default="Balanced",
-    help="How Controlnet injection should be prioritized.",
-)
-
 p.add_argument(
     "--use_lora",
     type=str,
     default="",
-    help="Use standalone LoRA weight using a HF ID or a checkpoint "
-    "file (~3 MB).",
+    help="Use standalone LoRA weight using a HF ID or a checkpoint file (~3 MB)",
 )
 
 p.add_argument(
     "--use_quantize",
     type=str,
     default="none",
-    help="Runs the quantized version of stable diffusion model. "
-    "This is currently in experimental phase. "
-    "Currently, only runs the stable-diffusion-2-1-base model in "
-    "int8 quantization.",
+    help="""Runs the quantized version of stable diffusion model. This is currently in experimental phase.
+            Currently, only runs the stable-diffusion-2-1-base model in int8 quantization.""",
 )
 
 p.add_argument(
     "--ondemand",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="Load and unload models for low VRAM.",
+    help="Load and unload models for low VRAM",
 )
 
-p.add_argument(
-    "--hf_auth_token",
-    type=str,
-    default=None,
-    help="Specify your own huggingface authentication tokens for models like Llama2.",
-)
-
-p.add_argument(
-    "--device_allocator_heap_key",
-    type=str,
-    default="",
-    help="Specify heap key for device caching allocator."
-    "Expected form: max_allocation_size;max_allocation_capacity;max_free_allocation_count"
-    "Example: --device_allocator_heap_key='*;1gib' (will limit caching on device to 1 gigabyte)",
-)
-
-p.add_argument(
-    "--autogen",
-    type=bool,
-    default="False",
-    help="Only used for a gradio workaround.",
-)
 ##############################################################################
-# IREE - Vulkan supported flags
+### IREE - Vulkan supported flags
 ##############################################################################
 
 p.add_argument(
     "--iree_vulkan_target_triple",
     type=str,
     default="",
-    help="Specify target triple for vulkan.",
+    help="Specify target triple for vulkan",
 )
 
 p.add_argument(
-    "--iree_metal_target_platform",
-    type=str,
-    default="",
-    help="Specify target triple for metal.",
+    "--vulkan_debug_utils",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="Profiles vulkan device and collects the .rdc info",
+)
+
+p.add_argument(
+    "--vulkan_large_heap_block_size",
+    default="2073741824",
+    help="flag for setting VMA preferredLargeHeapBlockSize for vulkan device, default is 4G",
+)
+
+p.add_argument(
+    "--vulkan_validation_layers",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="flag for disabling vulkan validation layers when benchmarking",
 )
 
 ##############################################################################
-# Misc. Debug and Optimization flags
+### Misc. Debug and Optimization flags
 ##############################################################################
 
 p.add_argument(
     "--use_compiled_scheduler",
     default=True,
     action=argparse.BooleanOptionalAction,
-    help="Use the default scheduler precompiled into the model if available.",
+    help="use the default scheduler precompiled into the model if available",
 )
 
 p.add_argument(
     "--local_tank_cache",
     default="",
-    help="Specify where to save downloaded shark_tank artifacts. "
-    "If this is not set, the default is ~/.local/shark_tank/.",
+    help="Specify where to save downloaded shark_tank artifacts. If this is not set, the default is ~/.local/shark_tank/.",
 )
 
 p.add_argument(
     "--dump_isa",
     default=False,
     action="store_true",
-    help="When enabled call amdllpc to get ISA dumps. "
-    "Use with dispatch benchmarks.",
+    help="When enabled call amdllpc to get ISA dumps. use with dispatch benchmarks.",
 )
 
 p.add_argument(
     "--dispatch_benchmarks",
     default=None,
-    help="Dispatches to return benchmark data on. "
-    'Use "All" for all, and None for none.',
+    help='dispatches to return benchamrk data on.  use "All" for all, and None for none.',
 )
 
 p.add_argument(
     "--dispatch_benchmarks_dir",
     default="temp_dispatch_benchmarks",
-    help="Directory where you want to store dispatch data "
-    'generated with "--dispatch_benchmarks".',
+    help='directory where you want to store dispatch data generated with "--dispatch_benchmarks"',
 )
 
 p.add_argument(
     "--enable_rgp",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="Flag for inserting debug frames between iterations "
-    "for use with rgp.",
+    help="flag for inserting debug frames between iterations for use with rgp.",
 )
 
 p.add_argument(
     "--hide_steps",
     default=True,
     action=argparse.BooleanOptionalAction,
-    help="Flag for hiding the details of iteration/sec for each step.",
+    help="flag for hiding the details of iteration/sec for each step.",
 )
 
 p.add_argument(
     "--warmup_count",
     type=int,
     default=0,
-    help="Flag setting warmup count for CLIP and VAE [>= 0].",
+    help="flag setting warmup count for clip and vae [>= 0].",
 )
 
 p.add_argument(
     "--clear_all",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="Flag to clear all mlir and vmfb from common locations. "
-    "Recompiling will take several minutes.",
+    help="flag to clear all mlir and vmfb from common locations. Recompiling will take several minutes",
 )
 
 p.add_argument(
     "--save_metadata_to_json",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="Flag for whether or not to save a generation information "
-    "json file with the image.",
+    help="flag for whether or not to save a generation information json file with the image.",
 )
 
 p.add_argument(
     "--write_metadata_to_png",
     default=True,
     action=argparse.BooleanOptionalAction,
-    help="Flag for whether or not to save generation information in "
-    "PNG chunk text to generated images.",
+    help="flag for whether or not to save generation information in PNG chunk text to generated images.",
 )
 
 p.add_argument(
     "--import_debug",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="If import_mlir is True, saves mlir via the debug option "
-    "in shark importer. Does nothing if import_mlir is false (the default).",
+    help="if import_mlir is True, saves mlir via the debug option in shark importer. Does nothing if import_mlir is false (the default)",
 )
-
-p.add_argument(
-    "--compile_debug",
-    default=False,
-    action=argparse.BooleanOptionalAction,
-    help="Flag to toggle debug assert/verify flags for imported IR in the"
-    "iree-compiler. Default to false.",
-)
-
-p.add_argument(
-    "--iree_constant_folding",
-    default=True,
-    action=argparse.BooleanOptionalAction,
-    help="Controls constant folding in iree-compile for all SD models.",
-)
-
-p.add_argument(
-    "--data_tiling",
-    default=False,
-    action=argparse.BooleanOptionalAction,
-    help="Controls data tiling in iree-compile for all SD models.",
-)
-
 ##############################################################################
-# Web UI flags
+### Web UI flags
 ##############################################################################
 
 p.add_argument(
     "--progress_bar",
     default=True,
     action=argparse.BooleanOptionalAction,
-    help="Flag for removing the progress bar animation during "
-    "image generation.",
+    help="flag for removing the progress bar animation during image generation",
 )
 
 p.add_argument(
     "--ckpt_dir",
     type=str,
     default="",
-    help="Path to directory where all .ckpts are stored in order to populate "
-    "them in the web UI.",
+    help="Path to directory where all .ckpts are stored in order to populate them in the web UI",
 )
 # TODO: replace API flag when these can be run together
 p.add_argument(
     "--ui",
     type=str,
     default="app" if os.name == "nt" else "web",
-    help="One of: [api, app, web].",
+    help="one of: [api, app, web]",
 )
 
 p.add_argument(
     "--share",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="Flag for generating a public URL.",
+    help="flag for generating a public URL",
 )
 
 p.add_argument(
     "--server_port",
     type=int,
     default=8080,
-    help="Flag for setting server port.",
+    help="flag for setting server port",
 )
 
 p.add_argument(
     "--api",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="Flag for enabling rest API.",
-)
-
-p.add_argument(
-    "--api_accept_origin",
-    action="append",
-    type=str,
-    help="An origin to be accepted by the REST api for Cross Origin"
-    "Resource Sharing (CORS). Use multiple times for multiple origins, "
-    'or use --api_accept_origin="*" to accept all origins. If no origins '
-    "are set no CORS headers will be returned by the api. Use, for "
-    "instance, if you need to access the REST api from Javascript running "
-    "in a web browser.",
-)
-
-p.add_argument(
-    "--debug",
-    default=False,
-    action=argparse.BooleanOptionalAction,
-    help="Flag for enabling debugging log in WebUI.",
+    help="flag for enabling rest API",
 )
 
 p.add_argument(
     "--output_gallery",
     default=True,
     action=argparse.BooleanOptionalAction,
-    help="Flag for removing the output gallery tab, and avoid exposing "
-    "images under --output_dir in the UI.",
+    help="flag for removing the output gallery tab, and avoid exposing images under --output_dir in the UI",
 )
 
 p.add_argument(
     "--output_gallery_followlinks",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="Flag for whether the output gallery tab in the UI should "
-    "follow symlinks when listing subdirectories under --output_dir.",
+    help="flag for whether the output gallery tab in the UI should follow symlinks when listing subdirectorys under --output_dir",
 )
 
 
 ##############################################################################
-# SD model auto-annotation flags
+### SD model auto-annotation flags
 ##############################################################################
 
 p.add_argument(
     "--annotation_output",
     type=path_expand,
     default="./",
-    help="Directory to save the annotated mlir file.",
+    help="Directory to save the annotated mlir file",
 )
 
 p.add_argument(
@@ -714,55 +566,33 @@ p.add_argument(
     "--save_annotation",
     default=False,
     action=argparse.BooleanOptionalAction,
-    help="Save annotated mlir file.",
+    help="Save annotated mlir file",
 )
 ##############################################################################
-# SD model auto-tuner flags
+### SD model auto-tuner flags
 ##############################################################################
 
 p.add_argument(
     "--tuned_config_dir",
     type=path_expand,
     default="./",
-    help="Directory to save the tuned config file.",
+    help="Directory to save the tuned config file",
 )
 
 p.add_argument(
     "--num_iters",
     type=int,
     default=400,
-    help="Number of iterations for tuning.",
+    help="Number of iterations for tuning",
 )
 
 p.add_argument(
     "--search_op",
     type=str,
     default="all",
-    help="Op to be optimized, options are matmul, bmm, conv and all.",
+    help="Op to be optimized, options are matmul, bmm, conv and all",
 )
 
-##############################################################################
-# DocuChat Flags
-##############################################################################
-
-p.add_argument(
-    "--run_docuchat_web",
-    default=False,
-    action=argparse.BooleanOptionalAction,
-    help="Specifies whether the docuchat's web version is running or not.",
-)
-
-##############################################################################
-# rocm Flags
-##############################################################################
-
-p.add_argument(
-    "--iree_rocm_target_chip",
-    type=str,
-    default="",
-    help="Add the rocm device architecture ex gfx1100, gfx90a, etc. Use `hipinfo` "
-    "or `iree-run-module --dump_devices=rocm` or `hipinfo` to get desired arch name",
-)
 
 args, unknown = p.parse_known_args()
 if args.import_debug:

apps/stable_diffusion/src/utils/stencils/__init__.py

@@ -1,3 +1,2 @@
 from apps.stable_diffusion.src.utils.stencils.canny import CannyDetector
 from apps.stable_diffusion.src.utils.stencils.openpose import OpenposeDetector
-from apps.stable_diffusion.src.utils.stencils.zoe import ZoeDetector

apps/stable_diffusion/src/utils/stencils/stencil_utils.py

@@ -1,44 +1,14 @@
 import numpy as np
 from PIL import Image
 import torch
-import os
-from pathlib import Path
-import torchvision
-import time
 from apps.stable_diffusion.src.utils.stencils import (
     CannyDetector,
     OpenposeDetector,
-    ZoeDetector,
 )
 
 stencil = {}
 
 
-def save_img(img):
-    from apps.stable_diffusion.src.utils import (
-        get_generated_imgs_path,
-        get_generated_imgs_todays_subdir,
-    )
-
-    subdir = Path(get_generated_imgs_path(), "preprocessed_control_hints")
-    os.makedirs(subdir, exist_ok=True)
-    if isinstance(img, Image.Image):
-        img.save(
-            os.path.join(
-                subdir, "controlnet_" + str(int(time.time())) + ".png"
-            )
-        )
-    elif isinstance(img, np.ndarray):
-        img = Image.fromarray(img)
-        img.save(os.path.join(subdir, str(int(time.time())) + ".png"))
-    else:
-        converter = torchvision.transforms.ToPILImage()
-        for i in img:
-            converter(i).save(
-                os.path.join(subdir, str(int(time.time())) + ".png")
-            )
-
-
 def HWC3(x):
     assert x.dtype == np.uint8
     if x.ndim == 2:
@@ -58,7 +28,7 @@ def HWC3(x):
         return y
 
 
-def controlnet_hint_reshaping(
+def controlnet_hint_shaping(
     controlnet_hint, height, width, dtype, num_images_per_prompt=1
 ):
     channels = 3
@@ -77,12 +47,10 @@ def controlnet_hint_reshaping(
                 )
             return controlnet_hint
         else:
-            return controlnet_hint_reshaping(
-                Image.fromarray(controlnet_hint.detach().numpy()),
-                height,
-                width,
-                dtype,
-                num_images_per_prompt,
+            raise ValueError(
+                f"Acceptble shape of `stencil` are any of ({channels}, {height}, {width}),"
+                + f" (1, {channels}, {height}, {width}) or ({num_images_per_prompt}, "
+                + f"{channels}, {height}, {width}) but is {controlnet_hint.shape}"
             )
     elif isinstance(controlnet_hint, np.ndarray):
         # np.ndarray: acceptable shape is any of hw, hwc, bhwc(b==1) or bhwc(b==num_images_per_promot)
@@ -109,38 +77,29 @@ def controlnet_hint_reshaping(
             )  # b h w c -> b c h w
             return controlnet_hint
         else:
-            return controlnet_hint_reshaping(
-                Image.fromarray(controlnet_hint),
-                height,
-                width,
-                dtype,
-                num_images_per_prompt,
+            raise ValueError(
+                f"Acceptble shape of `stencil` are any of ({width}, {channels}), "
+                + f"({height}, {width}, {channels}), "
+                + f"(1, {height}, {width}, {channels}) or "
+                + f"({num_images_per_prompt}, {channels}, {height}, {width}) but is {controlnet_hint.shape}"
             )
-
     elif isinstance(controlnet_hint, Image.Image):
-        controlnet_hint = controlnet_hint.convert(
-            "RGB"
-        )  # make sure 3 channel RGB format
         if controlnet_hint.size == (width, height):
-            controlnet_hint = np.array(controlnet_hint).astype(
-                np.float16
-            )  # to numpy
+            controlnet_hint = controlnet_hint.convert(
+                "RGB"
+            )  # make sure 3 channel RGB format
+            controlnet_hint = np.array(controlnet_hint)  # to numpy
             controlnet_hint = controlnet_hint[:, :, ::-1]  # RGB -> BGR
-            return controlnet_hint_reshaping(
-                controlnet_hint, height, width, dtype, num_images_per_prompt
+            return controlnet_hint_shaping(
+                controlnet_hint, height, width, num_images_per_prompt
             )
         else:
-            (hint_w, hint_h) = controlnet_hint.size
-            left = int((hint_w - width) / 2)
-            right = left + height
-            controlnet_hint = controlnet_hint.crop((left, 0, right, hint_h))
-            controlnet_hint = controlnet_hint.resize((width, height))
-        return controlnet_hint_reshaping(
-            controlnet_hint, height, width, dtype, num_images_per_prompt
-        )
+            raise ValueError(
+                f"Acceptable image size of `stencil` is ({width}, {height}) but is {controlnet_hint.size}"
+            )
     else:
         raise ValueError(
-            f"Acceptible controlnet input types are any of torch.Tensor, np.ndarray, PIL.Image.Image but is {type(controlnet_hint)}"
+            f"Acceptable type of `stencil` are any of torch.Tensor, np.ndarray, PIL.Image.Image but is {type(controlnet_hint)}"
         )
 
 
@@ -150,26 +109,17 @@ def controlnet_hint_conversion(
     controlnet_hint = None
     match use_stencil:
         case "canny":
-            print(
-                "Converting controlnet hint to edge detection mask with canny preprocessor."
-            )
+            print("Detecting edge with canny")
             controlnet_hint = hint_canny(image)
         case "openpose":
-            print(
-                "Detecting human pose in controlnet hint with openpose preprocessor."
-            )
+            print("Detecting human pose")
             controlnet_hint = hint_openpose(image)
         case "scribble":
-            print("Using your scribble as a controlnet hint.")
+            print("Working with scribble")
             controlnet_hint = hint_scribble(image)
-        case "zoedepth":
-            print(
-                "Converting controlnet hint to a depth mapping with ZoeDepth."
-            )
-            controlnet_hint = hint_zoedepth(image)
         case _:
             return None
-    controlnet_hint = controlnet_hint_reshaping(
+    controlnet_hint = controlnet_hint_shaping(
         controlnet_hint, height, width, dtype, num_images_per_prompt
     )
     return controlnet_hint
@@ -177,7 +127,7 @@ def controlnet_hint_conversion(
 
 stencil_to_model_id_map = {
     "canny": "lllyasviel/control_v11p_sd15_canny",
-    "zoedepth": "lllyasviel/control_v11f1p_sd15_depth",
+    "depth": "lllyasviel/control_v11p_sd15_depth",
     "hed": "lllyasviel/sd-controlnet-hed",
     "mlsd": "lllyasviel/control_v11p_sd15_mlsd",
     "normal": "lllyasviel/control_v11p_sd15_normalbae",
@@ -207,7 +157,6 @@ def hint_canny(
         detected_map = stencil["canny"](
             input_image, low_threshold, high_threshold
         )
-        save_img(detected_map)
         detected_map = HWC3(detected_map)
         return detected_map
 
@@ -223,7 +172,6 @@ def hint_openpose(
             stencil["openpose"] = OpenposeDetector()
 
         detected_map, _ = stencil["openpose"](input_image)
-        save_img(detected_map)
         detected_map = HWC3(detected_map)
         return detected_map
 
@@ -235,19 +183,4 @@ def hint_scribble(image: Image.Image):
 
         detected_map = np.zeros_like(input_image, dtype=np.uint8)
         detected_map[np.min(input_image, axis=2) < 127] = 255
-        save_img(detected_map)
-        return detected_map
-
-
-# Stencil 4. Depth (Only Zoe Preprocessing)
-def hint_zoedepth(image: Image.Image):
-    with torch.no_grad():
-        input_image = np.array(image)
-
-        if not "depth" in stencil:
-            stencil["depth"] = ZoeDetector()
-
-        detected_map = stencil["depth"](input_image)
-        save_img(detected_map)
-        detected_map = HWC3(detected_map)
         return detected_map

apps/stable_diffusion/src/utils/stencils/zoe/__init__.py

@@ -1,64 +0,0 @@
-import numpy as np
-import torch
-from pathlib import Path
-import requests
-
-
-from einops import rearrange
-
-remote_model_path = (
-    "https://huggingface.co/lllyasviel/Annotators/resolve/main/ZoeD_M12_N.pt"
-)
-
-
-class ZoeDetector:
-    def __init__(self):
-        cwd = Path.cwd()
-        ckpt_path = Path(cwd, "stencil_annotator")
-        ckpt_path.mkdir(parents=True, exist_ok=True)
-        modelpath = ckpt_path / "ZoeD_M12_N.pt"
-
-        with requests.get(remote_model_path, stream=True) as r:
-            r.raise_for_status()
-            with open(modelpath, "wb") as f:
-                for chunk in r.iter_content(chunk_size=8192):
-                    f.write(chunk)
-
-        model = torch.hub.load(
-            "monorimet/ZoeDepth:torch_update",
-            "ZoeD_N",
-            pretrained=False,
-            force_reload=False,
-        )
-
-        # Hack to fix the ZoeDepth import issue
-        model_keys = model.state_dict().keys()
-        loaded_dict = torch.load(modelpath, map_location=model.device)["model"]
-        loaded_keys = loaded_dict.keys()
-        for key in loaded_keys - model_keys:
-            loaded_dict.pop(key)
-
-        model.load_state_dict(loaded_dict)
-        model.eval()
-        self.model = model
-
-    def __call__(self, input_image):
-        assert input_image.ndim == 3
-        image_depth = input_image
-        with torch.no_grad():
-            image_depth = torch.from_numpy(image_depth).float()
-            image_depth = image_depth / 255.0
-            image_depth = rearrange(image_depth, "h w c -> 1 c h w")
-            depth = self.model.infer(image_depth)
-
-            depth = depth[0, 0].cpu().numpy()
-
-            vmin = np.percentile(depth, 2)
-            vmax = np.percentile(depth, 85)
-
-            depth -= vmin
-            depth /= vmax - vmin
-            depth = 1.0 - depth
-            depth_image = (depth * 255.0).clip(0, 255).astype(np.uint8)
-
-            return depth_image

apps/stable_diffusion/src/utils/utils.py

@@ -8,24 +8,17 @@ from datetime import datetime as dt
 from csv import DictWriter
 from pathlib import Path
 import numpy as np
-from random import (
-    randint,
-    seed as seed_random,
-    getstate as random_getstate,
-    setstate as random_setstate,
-)
+from random import randint
 import tempfile
 import torch
 from safetensors.torch import load_file
 from shark.shark_inference import SharkInference
-from shark.shark_importer import import_with_fx, save_mlir
+from shark.shark_importer import import_with_fx
 from shark.iree_utils.vulkan_utils import (
     set_iree_vulkan_runtime_flags,
     get_vulkan_target_triple,
-    get_iree_vulkan_runtime_flags,
 )
-from shark.iree_utils.metal_utils import get_metal_target_triple
-from shark.iree_utils.gpu_utils import get_cuda_sm_cc, get_iree_rocm_args
+from shark.iree_utils.gpu_utils import get_cuda_sm_cc
 from apps.stable_diffusion.src.utils.stable_args import args
 from apps.stable_diffusion.src.utils.resources import opt_flags
 from apps.stable_diffusion.src.utils.sd_annotation import sd_model_annotation
@@ -38,7 +31,6 @@ from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
 import requests
 from io import BytesIO
 from omegaconf import OmegaConf
-from cpuinfo import get_cpu_info
 
 
 def get_extended_name(model_name):
@@ -78,7 +70,7 @@ def _compile_module(shark_module, model_name, extra_args=[]):
                     )
                 )
             path = shark_module.save_module(
-                os.getcwd(), model_name, extra_args, debug=args.compile_debug
+                os.getcwd(), model_name, extra_args
             )
             shark_module.load_module(path, extra_args=extra_args)
     else:
@@ -87,9 +79,7 @@ def _compile_module(shark_module, model_name, extra_args=[]):
 
 
 # Downloads the model from shark_tank and returns the shark_module.
-def get_shark_model(tank_url, model_name, extra_args=None):
-    if extra_args is None:
-        extra_args = []
+def get_shark_model(tank_url, model_name, extra_args=[]):
     from shark.parser import shark_args
 
     # Set local shark_tank cache directory.
@@ -118,18 +108,16 @@ def compile_through_fx(
     is_f16=False,
     f16_input_mask=None,
     use_tuned=False,
-    save_dir="",
+    save_dir=tempfile.gettempdir(),
     debug=False,
     generate_vmfb=True,
-    extra_args=None,
+    extra_args=[],
     base_model_id=None,
     model_name=None,
     precision=None,
     return_mlir=False,
     device=None,
 ):
-    if extra_args is None:
-        extra_args = []
     if not return_mlir and model_name is not None:
         vmfb_path = get_vmfb_path_name(extended_model_name)
         if os.path.isfile(vmfb_path):
@@ -154,8 +142,8 @@ def compile_through_fx(
         f16_input_mask=f16_input_mask,
         debug=debug,
         model_name=extended_model_name,
+        save_dir=save_dir,
     )
-
     if use_tuned:
         if "vae" in extended_model_name.split("_")[0]:
             args.annotation_model = "vae"
@@ -168,14 +156,6 @@ def compile_through_fx(
             mlir_module, extended_model_name, base_model_id
         )
 
-    if not os.path.isdir(save_dir):
-        save_dir = ""
-
-    mlir_module = save_mlir(
-        mlir_module,
-        model_name=extended_model_name,
-        dir=save_dir,
-    )
     shark_module = SharkInference(
         mlir_module,
         device=args.device if device is None else device,
@@ -187,22 +167,20 @@ def compile_through_fx(
             mlir_module,
         )
 
+    del mlir_module
     gc.collect()
 
 
 def set_iree_runtime_flags():
-    # TODO: This function should be device-agnostic and piped properly
-    # to general runtime driver init.
-    vulkan_runtime_flags = get_iree_vulkan_runtime_flags()
+    vulkan_runtime_flags = [
+        f"--vulkan_large_heap_block_size={args.vulkan_large_heap_block_size}",
+        f"--vulkan_validation_layers={'true' if args.vulkan_validation_layers else 'false'}",
+    ]
     if args.enable_rgp:
         vulkan_runtime_flags += [
             f"--enable_rgp=true",
             f"--vulkan_debug_utils=true",
         ]
-    if args.device_allocator_heap_key:
-        vulkan_runtime_flags += [
-            f"--device_allocator=caching:device_local={args.device_allocator_heap_key}",
-        ]
     set_iree_vulkan_runtime_flags(flags=vulkan_runtime_flags)
 
 
@@ -225,15 +203,13 @@ def get_device_mapping(driver, key_combination=3):
     specific devices for execution
     Args:
         driver (str): execution driver (vulkan, cuda, rocm, etc)
-        key_combination (int, optional): choice for mapping value for
-            device name.
+        key_combination (int, optional): choice for mapping value for device name.
         1 : path
         2 : name
         3 : (name, path)
         Defaults to 3.
     Returns:
-        dict: map to possible device names user can input mapped to desired
-            combination of name/path.
+        dict: map to possible device names user can input mapped to desired combination of name/path.
     """
     from shark.iree_utils._common import iree_device_map
 
@@ -247,7 +223,7 @@ def get_device_mapping(driver, key_combination=3):
         if key_combination == 2:
             return dev_dict["name"]
         if key_combination == 3:
-            return dev_dict["name"], f"{driver}://{dev_dict['path']}"
+            return (dev_dict["name"], f"{driver}://{dev_dict['path']}")
 
     # mapping driver name to default device (driver://0)
     device_map[f"{driver}"] = get_output_value(device_list[0])
@@ -260,12 +236,10 @@ def get_device_mapping(driver, key_combination=3):
 
 
 def map_device_to_name_path(device, key_combination=3):
-    """Gives the appropriate device data (supported name/path) for user
-        selected execution device
+    """Gives the appropriate device data (supported name/path) for user selected execution device
     Args:
         device (str): user
-        key_combination (int, optional): choice for mapping value for
-            device name.
+        key_combination (int, optional): choice for mapping value for device name.
         1 : path
         2 : name
         3 : (name, path)
@@ -273,8 +247,7 @@ def map_device_to_name_path(device, key_combination=3):
     Raises:
         ValueError:
     Returns:
-        str / tuple: returns the mapping str or tuple of mapping str for
-        the device depending on key_combination value
+        str / tuple: returns the mapping str or tuple of mapping str for the device depending on key_combination value
     """
     driver = device.split("://")[0]
     device_map = get_device_mapping(driver, key_combination)
@@ -297,21 +270,10 @@ def set_init_device_flags():
             if triple is not None:
                 args.iree_vulkan_target_triple = triple
         print(
-            f"Found device {device_name}. Using target triple "
-            f"{args.iree_vulkan_target_triple}."
+            f"Found device {device_name}. Using target triple {args.iree_vulkan_target_triple}."
         )
     elif "cuda" in args.device:
         args.device = "cuda"
-    elif "metal" in args.device:
-        device_name, args.device = map_device_to_name_path(args.device)
-        if not args.iree_metal_target_platform:
-            triple = get_metal_target_triple(device_name)
-            if triple is not None:
-                args.iree_metal_target_platform = triple.split("-")[-1]
-        print(
-            f"Found device {device_name}. Using target triple "
-            f"{args.iree_metal_target_platform}."
-        )
     elif "cpu" in args.device:
         args.device = "cpu"
 
@@ -343,17 +305,6 @@ def set_init_device_flags():
     ):
         args.use_tuned = False
 
-    elif (
-        args.height != args.width
-        and "rdna2" in args.iree_vulkan_target_triple
-        and base_model_id
-        not in [
-            "CompVis/stable-diffusion-v1-4",
-            "runwayml/stable-diffusion-v1-5",
-        ]
-    ):
-        args.use_tuned = False
-
     elif base_model_id not in [
         "Linaqruf/anything-v3.0",
         "dreamlike-art/dreamlike-diffusion-1.0",
@@ -408,8 +359,7 @@ def set_init_device_flags():
 
     if args.use_tuned:
         print(
-            f"Using tuned models for {base_model_id}(fp16) on "
-            f"device {args.device}."
+            f"Using tuned models for {base_model_id}(fp16) on device {args.device}."
         )
     else:
         print("Tuned models are currently not supported for this setting.")
@@ -467,52 +417,18 @@ def get_available_devices():
         except:
             print(f"{driver_name} devices are not available.")
         else:
-            cpu_name = get_cpu_info()["brand_raw"]
             for i, device in enumerate(device_list_dict):
-                device_name = (
-                    cpu_name if device["name"] == "default" else device["name"]
-                )
-                if "local" in driver_name:
-                    device_list.append(
-                        f"{device_name} => {driver_name.replace('local', 'cpu')}"
-                    )
-                else:
-                    # for drivers with single devices
-                    # let the default device be selected without any indexing
-                    if len(device_list_dict) == 1:
-                        device_list.append(f"{device_name} => {driver_name}")
-                    else:
-                        device_list.append(
-                            f"{device_name} => {driver_name}://{i}"
-                        )
+                device_list.append(f"{device['name']} => {driver_name}://{i}")
         return device_list
 
     set_iree_runtime_flags()
 
     available_devices = []
-    from shark.iree_utils.vulkan_utils import (
-        get_all_vulkan_devices,
-    )
-
-    vulkaninfo_list = get_all_vulkan_devices()
-    vulkan_devices = []
-    id = 0
-    for device in vulkaninfo_list:
-        vulkan_devices.append(f"{device.strip()} => vulkan://{id}")
-        id += 1
-    if id != 0:
-        print(f"vulkan devices are available.")
+    vulkan_devices = get_devices_by_name("vulkan")
     available_devices.extend(vulkan_devices)
-    metal_devices = get_devices_by_name("metal")
-    available_devices.extend(metal_devices)
     cuda_devices = get_devices_by_name("cuda")
     available_devices.extend(cuda_devices)
-    rocm_devices = get_devices_by_name("rocm")
-    available_devices.extend(rocm_devices)
-    cpu_device = get_devices_by_name("cpu-sync")
-    available_devices.extend(cpu_device)
-    cpu_device = get_devices_by_name("cpu-task")
-    available_devices.extend(cpu_device)
+    available_devices.append("device => cpu")
     return available_devices
 
 
@@ -532,17 +448,10 @@ def get_opt_flags(model, precision="fp16"):
         iree_flags.append(
             f"-iree-vulkan-target-triple={args.iree_vulkan_target_triple}"
         )
-    if "rocm" in args.device:
-        rocm_args = get_iree_rocm_args()
-        iree_flags.extend(rocm_args)
-        print(iree_flags)
-    if args.iree_constant_folding == False:
-        iree_flags.append("--iree-opt-const-expr-hoisting=False")
-        iree_flags.append(
-            "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807"
-        )
-    if args.data_tiling == False:
-        iree_flags.append("--iree-opt-data-tiling=False")
+
+    # Disable bindings fusion to work with moltenVK.
+    if sys.platform == "darwin":
+        iree_flags.append("-iree-stream-fuse-binding=false")
 
     if "default_compilation_flags" in opt_flags[model][is_tuned][precision]:
         iree_flags += opt_flags[model][is_tuned][precision][
@@ -565,10 +474,6 @@ def get_opt_flags(model, precision="fp16"):
         iree_flags += opt_flags[model][is_tuned][precision][
             "specified_compilation_flags"
         ][device]
-    if "vae" not in model:
-        # Due to lack of support for multi-reduce, we always collapse reduction
-        # dims before dispatch formation right now.
-        iree_flags += ["--iree-flow-collapse-reduction-dims"]
     return iree_flags
 
 
@@ -600,17 +505,17 @@ def preprocessCKPT(custom_weights, is_inpaint=False):
     from_safetensors = (
         True if custom_weights.lower().endswith(".safetensors") else False
     )
-    # EMA weights usually yield higher quality images for inference but
-    # non-EMA weights have been yielding better results in our case.
-    # TODO: Add an option `--ema` (`--no-ema`) for users to specify if
-    #  they want to go for EMA weight extraction or not.
+    # EMA weights usually yield higher quality images for inference but non-EMA weights have
+    # been yielding better results in our case.
+    # TODO: Add an option `--ema` (`--no-ema`) for users to specify if they want to go for EMA
+    #       weight extraction or not.
     extract_ema = False
     print(
         "Loading diffusers' pipeline from original stable diffusion checkpoint"
     )
     num_in_channels = 9 if is_inpaint else 4
     pipe = download_from_original_stable_diffusion_ckpt(
-        checkpoint_path_or_dict=custom_weights,
+        checkpoint_path=custom_weights,
         extract_ema=extract_ema,
         from_safetensors=from_safetensors,
         num_in_channels=num_in_channels,
@@ -624,10 +529,7 @@ def convert_original_vae(vae_checkpoint):
     for key in list(vae_checkpoint.keys()):
         vae_state_dict["first_stage_model." + key] = vae_checkpoint.get(key)
 
-    config_url = (
-        "https://raw.githubusercontent.com/CompVis/stable-diffusion/"
-        "main/configs/stable-diffusion/v1-inference.yaml"
-    )
+    config_url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"
     original_config_file = BytesIO(requests.get(config_url).content)
     original_config = OmegaConf.load(original_config_file)
     vae_config = create_vae_diffusers_config(original_config, image_size=512)
@@ -748,7 +650,7 @@ def update_lora_weight(model, use_lora, model_name):
 
 
 # `fetch_and_update_base_model_id` is a resource utility function which
-# helps to maintain mapping of the model to run with its base model.
+# helps maintaining mapping of the model to run with its base model.
 # If `base_model` is "", then this function tries to fetch the base model
 # info for the `model_to_run`.
 def fetch_and_update_base_model_id(model_to_run, base_model=""):
@@ -765,17 +667,14 @@ def fetch_and_update_base_model_id(model_to_run, base_model=""):
                 return base_model
     elif base_model == "":
         return base_model
-    # Update JSON data to contain an entry mapping model_to_run with
-    # base_model.
+    # Update JSON data to contain an entry mapping model_to_run with base_model.
     json_data.update(data)
     with open(variants_path, "w", encoding="utf-8") as jsonFile:
         json.dump(json_data, jsonFile)
 
 
-# Generate and return a new seed if the provided one is not in the
-# supported range (including -1)
-def sanitize_seed(seed: int | str):
-    seed = int(seed)
+# Generate and return a new seed if the provided one is not in the supported range (including -1)
+def sanitize_seed(seed):
     uint32_info = np.iinfo(np.uint32)
     uint32_min, uint32_max = uint32_info.min, uint32_info.max
     if seed < uint32_min or seed >= uint32_max:
@@ -783,57 +682,6 @@ def sanitize_seed(seed: int | str):
     return seed
 
 
-# take a seed expression in an input format and convert it to
-# a list of integers, where possible
-def parse_seed_input(seed_input: str | list | int):
-    if isinstance(seed_input, str):
-        try:
-            seed_input = json.loads(seed_input)
-        except (ValueError, TypeError):
-            seed_input = None
-
-    if isinstance(seed_input, int):
-        return [seed_input]
-
-    if isinstance(seed_input, list) and all(
-        type(seed) is int for seed in seed_input
-    ):
-        return seed_input
-
-    raise TypeError(
-        "Seed input must be an integer or an array of integers in JSON format"
-    )
-
-
-# Generate a set of seeds from an input expression for batch_count batches,
-# optionally using that input as the rng seed for any randomly generated seeds.
-def batch_seeds(
-    seed_input: str | list | int, batch_count: int, repeatable=False
-):
-    # turn the input into a list if possible
-    seeds = parse_seed_input(seed_input)
-
-    # slice or pad the list to be of batch_count length
-    seeds = seeds[:batch_count] + [-1] * (batch_count - len(seeds))
-
-    if repeatable:
-        if all(seed < 0 for seed in seeds):
-            seeds[0] = sanitize_seed(seeds[0])
-
-        # set seed for the rng based on what we have so far
-        saved_random_state = random_getstate()
-        seed_random(str([n for n in seeds if n > -1]))
-
-    # generate any seeds that are unspecified
-    seeds = [sanitize_seed(seed) for seed in seeds]
-
-    if repeatable:
-        # reset the rng back to normal
-        random_setstate(saved_random_state)
-
-    return seeds
-
-
 # clear all the cached objects to recompile cleanly.
 def clear_all():
     print("CLEARING ALL, EXPECT SEVERAL MINUTES TO RECOMPILE")
@@ -844,8 +692,7 @@ def clear_all():
     for vmfb in vmfbs:
         if os.path.exists(vmfb):
             os.remove(vmfb)
-    # Temporary workaround of deleting yaml files to incorporate
-    # diffusers' pipeline.
+    # Temporary workaround of deleting yaml files to incorporate diffusers' pipeline.
     # TODO: Remove this once we have better weight updation logic.
     inference_yaml = ["v2-inference-v.yaml", "v1-inference.yaml"]
     for yaml in inference_yaml:
@@ -861,8 +708,6 @@ def clear_all():
     elif os.name == "unix":
         shutil.rmtree(os.path.join(home, ".cache/AMD/VkCache"))
         shutil.rmtree(os.path.join(home, ".local/shark_tank"))
-    if args.local_tank_cache != "":
-        shutil.rmtree(args.local_tank_cache)
 
 
 def get_generated_imgs_path() -> Path:
@@ -876,9 +721,7 @@ def get_generated_imgs_todays_subdir() -> str:
 
 
 # save output images and the inputs corresponding to it.
-def save_output_img(output_img, img_seed, extra_info=None):
-    if extra_info is None:
-        extra_info = {}
+def save_output_img(output_img, img_seed, extra_info={}):
     generated_imgs_path = Path(
         get_generated_imgs_path(), get_generated_imgs_todays_subdir()
     )
@@ -886,20 +729,14 @@ def save_output_img(output_img, img_seed, extra_info=None):
     csv_path = Path(generated_imgs_path, "imgs_details.csv")
 
     prompt_slice = re.sub("[^a-zA-Z0-9]", "_", args.prompts[0][:15])
-    out_img_name = f"{dt.now().strftime('%H%M%S')}_{prompt_slice}_{img_seed}"
+    out_img_name = (
+        f"{prompt_slice}_{img_seed}_{dt.now().strftime('%y%m%d_%H%M%S')}"
+    )
 
     img_model = args.hf_model_id
     if args.ckpt_loc:
         img_model = Path(os.path.basename(args.ckpt_loc)).stem
 
-    img_vae = None
-    if args.custom_vae:
-        img_vae = Path(os.path.basename(args.custom_vae)).stem
-
-    img_lora = None
-    if args.use_lora:
-        img_lora = Path(os.path.basename(args.use_lora)).stem
-
     if args.output_img_format == "jpg":
         out_img_path = Path(generated_imgs_path, f"{out_img_name}.jpg")
         output_img.save(out_img_path, quality=95, subsampling=0)
@@ -908,39 +745,19 @@ def save_output_img(output_img, img_seed, extra_info=None):
         pngInfo = PngImagePlugin.PngInfo()
 
         if args.write_metadata_to_png:
-            # Using a conditional expression caused problems, so setting a new
-            # variable for now.
-            if args.use_hiresfix:
-                png_size_text = f"{args.hiresfix_width}x{args.hiresfix_height}"
-            else:
-                png_size_text = f"{args.width}x{args.height}"
-
             pngInfo.add_text(
                 "parameters",
-                f"{args.prompts[0]}"
-                f"\nNegative prompt: {args.negative_prompts[0]}"
-                f"\nSteps: {args.steps},"
-                f"Sampler: {args.scheduler}, "
-                f"CFG scale: {args.guidance_scale}, "
-                f"Seed: {img_seed},"
-                f"Size: {png_size_text}, "
-                f"Model: {img_model}, "
-                f"VAE: {img_vae}, "
-                f"LoRA: {img_lora}",
+                f"{args.prompts[0]}\nNegative prompt: {args.negative_prompts[0]}\nSteps:{args.steps}, Sampler: {args.scheduler}, CFG scale: {args.guidance_scale}, Seed: {img_seed}, Size: {args.width}x{args.height}, Model: {img_model}",
             )
 
         output_img.save(out_img_path, "PNG", pnginfo=pngInfo)
 
         if args.output_img_format not in ["png", "jpg"]:
             print(
-                f"[ERROR] Format {args.output_img_format} is not "
-                f"supported yet. Image saved as png instead."
-                f"Supported formats: png / jpg"
+                f"[ERROR] Format {args.output_img_format} is not supported yet."
+                "Image saved as png instead. Supported formats: png / jpg"
             )
 
-    # To be as low-impact as possible to the existing CSV format, we append
-    # "VAE" and "LORA" to the end. However, it does not fit the hierarchy of
-    # importance for each data point. Something to consider.
     new_entry = {
         "VARIANT": img_model,
         "SCHEDULER": args.scheduler,
@@ -950,23 +767,16 @@ def save_output_img(output_img, img_seed, extra_info=None):
         "CFG_SCALE": args.guidance_scale,
         "PRECISION": args.precision,
         "STEPS": args.steps,
-        "HEIGHT": args.height
-        if not args.use_hiresfix
-        else args.hiresfix_height,
-        "WIDTH": args.width if not args.use_hiresfix else args.hiresfix_width,
+        "HEIGHT": args.height,
+        "WIDTH": args.width,
         "MAX_LENGTH": args.max_length,
         "OUTPUT": out_img_path,
-        "VAE": img_vae,
-        "LORA": img_lora,
     }
 
     new_entry.update(extra_info)
 
-    csv_mode = "a" if os.path.isfile(csv_path) else "w"
-    with open(csv_path, csv_mode, encoding="utf-8") as csv_obj:
+    with open(csv_path, "a", encoding="utf-8") as csv_obj:
         dictwriter_obj = DictWriter(csv_obj, fieldnames=list(new_entry.keys()))
-        if csv_mode == "w":
-            dictwriter_obj.writeheader()
         dictwriter_obj.writerow(new_entry)
         csv_obj.close()
 
@@ -980,35 +790,21 @@ def save_output_img(output_img, img_seed, extra_info=None):
 def get_generation_text_info(seeds, device):
     text_output = f"prompt={args.prompts}"
     text_output += f"\nnegative prompt={args.negative_prompts}"
-    text_output += (
-        f"\nmodel_id={args.hf_model_id}, " f"ckpt_loc={args.ckpt_loc}"
-    )
-    text_output += f"\nscheduler={args.scheduler}, " f"device={device}"
-    text_output += (
-        f"\nsteps={args.steps}, "
-        f"guidance_scale={args.guidance_scale}, "
-        f"seed={seeds}"
-    )
-    text_output += (
-        f"\nsize={args.height}x{args.width}, "
-        if not args.use_hiresfix
-        else f"\nsize={args.hiresfix_height}x{args.hiresfix_width}, "
-    )
-    text_output += (
-        f"batch_count={args.batch_count}, "
-        f"batch_size={args.batch_size}, "
-        f"max_length={args.max_length}"
-    )
+    text_output += f"\nmodel_id={args.hf_model_id}, ckpt_loc={args.ckpt_loc}"
+    text_output += f"\nscheduler={args.scheduler}, device={device}"
+    text_output += f"\nsteps={args.steps}, guidance_scale={args.guidance_scale}, seed={seeds}"
+    text_output += f"\nsize={args.height}x{args.width}, batch_count={args.batch_count}, batch_size={args.batch_size}, max_length={args.max_length}"
 
     return text_output
 
 
-# For stencil, the input image can be of any size, but we need to ensure that
-# it conforms with our model constraints :-
+# For stencil, the input image can be of any size but we need to ensure that
+# it conforms with our model contraints :-
 #   Both width and height should be in the range of [128, 768] and multiple of 8.
 # This utility function performs the transformation on the input image while
 # also maintaining the aspect ratio before sending it to the stencil pipeline.
-def resize_stencil(image: Image.Image, width, height):
+def resize_stencil(image: Image.Image):
+    width, height = image.size
     aspect_ratio = width / height
     min_size = min(width, height)
     if min_size < 128:

apps/stable_diffusion/studio_bundle.spec

@@ -1,54 +0,0 @@
-# -*- mode: python ; coding: utf-8 -*-
-from apps.stable_diffusion.shark_studio_imports import pathex, datas, hiddenimports
-
-binaries = []
-
-block_cipher = None
-
-a = Analysis(
-    ['web\\index.py'],
-    pathex=pathex,
-    binaries=binaries,
-    datas=datas,
-    hiddenimports=hiddenimports,
-    hookspath=[],
-    hooksconfig={},
-    runtime_hooks=[],
-    excludes=[],
-    win_no_prefer_redirects=False,
-    win_private_assemblies=False,
-    cipher=block_cipher,
-    noarchive=False,
-    module_collection_mode={
-        'gradio': 'py',  # Collect gradio package as source .py files
-    },
-)
-pyz = PYZ(a.pure, a.zipped_data, cipher=block_cipher)
-
-exe = EXE(
-    pyz,
-    a.scripts,
-    [],
-    exclude_binaries=True,
-    name='studio_bundle',
-    debug=False,
-    bootloader_ignore_signals=False,
-    strip=False,
-    upx=True,
-    console=True,
-    disable_windowed_traceback=False,
-    argv_emulation=False,
-    target_arch=None,
-    codesign_identity=None,
-    entitlements_file=None,
-)
-coll = COLLECT(
-    exe,
-    a.binaries,
-    a.zipfiles,
-    a.datas,
-    strip=False,
-    upx=True,
-    upx_exclude=[],
-    name='studio_bundle',
-)

apps/stable_diffusion/web/api/__init__.py

@@ -1 +0,0 @@
-from apps.stable_diffusion.web.api.sdapi_v1 import sdapi

apps/stable_diffusion/web/api/sdapi_v1.py

@@ -1,579 +0,0 @@
-import os
-
-from collections import defaultdict
-from enum import Enum
-from fastapi import FastAPI
-from pydantic import BaseModel, Field, conlist, model_validator
-
-from apps.stable_diffusion.web.api.utils import (
-    frozen_args,
-    sampler_aliases,
-    encode_pil_to_base64,
-    decode_base64_to_image,
-    get_model_from_request,
-    get_scheduler_from_request,
-    get_lora_params,
-    get_device,
-    GenerationInputData,
-    GenerationResponseData,
-)
-
-from apps.stable_diffusion.web.ui.utils import (
-    get_custom_model_files,
-    get_custom_model_pathfile,
-    predefined_models,
-    predefined_paint_models,
-    predefined_upscaler_models,
-    scheduler_list,
-)
-from apps.stable_diffusion.web.ui.txt2img_ui import txt2img_inf
-from apps.stable_diffusion.web.ui.img2img_ui import img2img_inf
-from apps.stable_diffusion.web.ui.inpaint_ui import inpaint_inf
-from apps.stable_diffusion.web.ui.outpaint_ui import outpaint_inf
-from apps.stable_diffusion.web.ui.upscaler_ui import upscaler_inf
-
-sdapi = FastAPI()
-
-
-# Rest API: /sdapi/v1/sd-models (lists available models)
-class AppParam(str, Enum):
-    txt2img = "txt2img"
-    img2img = "img2img"
-    inpaint = "inpaint"
-    outpaint = "outpaint"
-    upscaler = "upscaler"
-
-
-@sdapi.get(
-    "/v1/sd-models",
-    summary="lists available models",
-    description=(
-        "This is all the models that this server currently knows about.\n "
-        "Models listed may still have a compilation and build pending that "
-        "will be triggered the first time they are used."
-    ),
-)
-def sd_models_api(app: AppParam = frozen_args.app):
-    match app:
-        case "inpaint" | "outpaint":
-            checkpoint_type = "inpainting"
-            predefined = predefined_paint_models
-        case "upscaler":
-            checkpoint_type = "upscaler"
-            predefined = predefined_upscaler_models
-        case _:
-            checkpoint_type = ""
-            predefined = predefined_models
-
-    return [
-        {
-            "title": model_file,
-            "model_name": model_file,
-            "hash": None,
-            "sha256": None,
-            "filename": get_custom_model_pathfile(model_file),
-            "config": None,
-        }
-        for model_file in get_custom_model_files(
-            custom_checkpoint_type=checkpoint_type
-        )
-    ] + [
-        {
-            "title": model,
-            "model_name": model,
-            "hash": None,
-            "sha256": None,
-            "filename": None,
-            "config": None,
-        }
-        for model in predefined
-    ]
-
-
-# Rest API: /sdapi/v1/samplers (lists schedulers)
-@sdapi.get(
-    "/v1/samplers",
-    summary="lists available schedulers/samplers",
-    description=(
-        "These are all the Schedulers defined and available. Not "
-        "every scheduler is compatible with all apis. Aliases are "
-        "equivalent samplers in A1111 if they are known."
-    ),
-)
-def sd_samplers_api():
-    reverse_sampler_aliases = defaultdict(list)
-    for key, value in sampler_aliases.items():
-        reverse_sampler_aliases[value].append(key)
-
-    return (
-        {
-            "name": scheduler,
-            "aliases": reverse_sampler_aliases.get(scheduler, []),
-            "options": {},
-        }
-        for scheduler in scheduler_list
-    )
-
-
-# Rest API: /sdapi/v1/options (lists application level options)
-@sdapi.get(
-    "/v1/options",
-    summary="lists current settings of application level options",
-    description=(
-        "A subset of the command line arguments set at startup renamed "
-        "to correspond to the A1111 naming. Only a small subset of A1111 "
-        "options are returned."
-    ),
-)
-def options_api():
-    # This is mostly just enough to support what Koboldcpp wants, with a
-    # few other things that seemed obvious
-    return {
-        "samples_save": True,
-        "samples_format": frozen_args.output_img_format,
-        "sd_model_checkpoint": os.path.basename(frozen_args.ckpt_loc)
-        if frozen_args.ckpt_loc
-        else frozen_args.hf_model_id,
-        "sd_lora": frozen_args.use_lora,
-        "sd_vae": frozen_args.custom_vae or "Automatic",
-        "enable_pnginfo": frozen_args.write_metadata_to_png,
-    }
-
-
-# Rest API: /sdapi/v1/cmd-flags (lists command line argument settings)
-@sdapi.get(
-    "/v1/cmd-flags",
-    summary="lists the command line arguments value that were set on startup.",
-)
-def cmd_flags_api():
-    return vars(frozen_args)
-
-
-# Rest API: /sdapi/v1/txt2img (Text to image)
-class ModelOverrideSettings(BaseModel):
-    sd_model_checkpoint: str = get_model_from_request(
-        fallback_model="stabilityai/stable-diffusion-2-1-base"
-    )
-
-
-class Txt2ImgInputData(GenerationInputData):
-    enable_hr: bool = frozen_args.use_hiresfix
-    hr_resize_y: int = Field(
-        default=frozen_args.hiresfix_height, ge=128, le=768, multiple_of=8
-    )
-    hr_resize_x: int = Field(
-        default=frozen_args.hiresfix_width, ge=128, le=768, multiple_of=8
-    )
-    override_settings: ModelOverrideSettings = None
-
-
-@sdapi.post(
-    "/v1/txt2img",
-    summary="Does text to image generation",
-    response_model=GenerationResponseData,
-)
-def txt2img_api(InputData: Txt2ImgInputData):
-    model_id = get_model_from_request(
-        InputData,
-        fallback_model="stabilityai/stable-diffusion-2-1-base",
-    )
-    scheduler = get_scheduler_from_request(
-        InputData, "txt2img_hires" if InputData.enable_hr else "txt2img"
-    )
-    (lora_weights, lora_hf_id) = get_lora_params(frozen_args.use_lora)
-
-    print(
-        f"Prompt: {InputData.prompt}, "
-        f"Negative Prompt: {InputData.negative_prompt}, "
-        f"Seed: {InputData.seed},"
-        f"Model: {model_id}, "
-        f"Scheduler: {scheduler}. "
-    )
-
-    res = txt2img_inf(
-        InputData.prompt,
-        InputData.negative_prompt,
-        InputData.height,
-        InputData.width,
-        InputData.steps,
-        InputData.cfg_scale,
-        InputData.seed,
-        batch_count=InputData.n_iter,
-        batch_size=1,
-        scheduler=scheduler,
-        model_id=model_id,
-        custom_vae=frozen_args.custom_vae or "None",
-        precision="fp16",
-        device=get_device(frozen_args.device),
-        max_length=frozen_args.max_length,
-        save_metadata_to_json=frozen_args.save_metadata_to_json,
-        save_metadata_to_png=frozen_args.write_metadata_to_png,
-        lora_weights=lora_weights,
-        lora_hf_id=lora_hf_id,
-        ondemand=frozen_args.ondemand,
-        repeatable_seeds=False,
-        use_hiresfix=InputData.enable_hr,
-        hiresfix_height=InputData.hr_resize_y,
-        hiresfix_width=InputData.hr_resize_x,
-        hiresfix_strength=frozen_args.hiresfix_strength,
-        resample_type=frozen_args.resample_type,
-    )
-
-    # Since we're not streaming we just want the last generator result
-    for items_so_far in res:
-        items = items_so_far
-
-    return {
-        "images": encode_pil_to_base64(items[0]),
-        "parameters": {},
-        "info": items[1],
-    }
-
-
-# Rest API: /sdapi/v1/img2img (Image to image)
-class StencilParam(str, Enum):
-    canny = "canny"
-    openpose = "openpose"
-    scribble = "scribble"
-    zoedepth = "zoedepth"
-
-
-class Img2ImgInputData(GenerationInputData):
-    init_images: conlist(str, min_length=1, max_length=2)
-    denoising_strength: float = frozen_args.strength
-    use_stencil: StencilParam = frozen_args.use_stencil
-    override_settings: ModelOverrideSettings = None
-
-    @model_validator(mode="after")
-    def check_image_supplied_for_scribble_stencil(self) -> "Img2ImgInputData":
-        if (
-            self.use_stencil == StencilParam.scribble
-            and len(self.init_images) < 2
-        ):
-            raise ValueError(
-                "a second image must be supplied for the controlnet:scribble stencil"
-            )
-
-        return self
-
-
-@sdapi.post(
-    "/v1/img2img",
-    summary="Does image to image generation",
-    response_model=GenerationResponseData,
-)
-def img2img_api(
-    InputData: Img2ImgInputData,
-):
-    model_id = get_model_from_request(
-        InputData,
-        fallback_model="stabilityai/stable-diffusion-2-1-base",
-    )
-    scheduler = get_scheduler_from_request(InputData, "img2img")
-    (lora_weights, lora_hf_id) = get_lora_params(frozen_args.use_lora)
-
-    init_image = decode_base64_to_image(InputData.init_images[0])
-    mask_image = (
-        decode_base64_to_image(InputData.init_images[1])
-        if len(InputData.init_images) > 1
-        else None
-    )
-
-    print(
-        f"Prompt: {InputData.prompt}, "
-        f"Negative Prompt: {InputData.negative_prompt}, "
-        f"Seed: {InputData.seed}, "
-        f"Model: {model_id}, "
-        f"Scheduler: {scheduler}."
-    )
-
-    res = img2img_inf(
-        InputData.prompt,
-        InputData.negative_prompt,
-        {"image": init_image, "mask": mask_image},
-        InputData.height,
-        InputData.width,
-        InputData.steps,
-        InputData.denoising_strength,
-        InputData.cfg_scale,
-        InputData.seed,
-        batch_count=InputData.n_iter,
-        batch_size=1,
-        scheduler=scheduler,
-        model_id=model_id,
-        custom_vae=frozen_args.custom_vae or "None",
-        precision="fp16",
-        device=get_device(frozen_args.device),
-        max_length=frozen_args.max_length,
-        use_stencil=InputData.use_stencil,
-        save_metadata_to_json=frozen_args.save_metadata_to_json,
-        save_metadata_to_png=frozen_args.write_metadata_to_png,
-        lora_weights=lora_weights,
-        lora_hf_id=lora_hf_id,
-        ondemand=frozen_args.ondemand,
-        repeatable_seeds=False,
-        resample_type=frozen_args.resample_type,
-    )
-
-    # Since we're not streaming we just want the last generator result
-    for items_so_far in res:
-        items = items_so_far
-
-    return {
-        "images": encode_pil_to_base64(items[0]),
-        "parameters": {},
-        "info": items[1],
-    }
-
-
-# Rest API: /sdapi/v1/inpaint (Inpainting)
-class PaintModelOverideSettings(BaseModel):
-    sd_model_checkpoint: str = get_model_from_request(
-        checkpoint_type="inpainting",
-        fallback_model="stabilityai/stable-diffusion-2-inpainting",
-    )
-
-
-class InpaintInputData(GenerationInputData):
-    image: str = Field(description="Base64 encoded input image")
-    mask: str = Field(description="Base64 encoded mask image")
-    is_full_res: bool = False  # Is this setting backwards in the UI?
-    full_res_padding: int = Field(default=32, ge=0, le=256, multiple_of=4)
-    denoising_strength: float = frozen_args.strength
-    use_stencil: StencilParam = frozen_args.use_stencil
-    override_settings: PaintModelOverideSettings = None
-
-
-@sdapi.post(
-    "/v1/inpaint",
-    summary="Does inpainting generation on an image",
-    response_model=GenerationResponseData,
-)
-def inpaint_api(
-    InputData: InpaintInputData,
-):
-    model_id = get_model_from_request(
-        InputData,
-        checkpoint_type="inpainting",
-        fallback_model="stabilityai/stable-diffusion-2-inpainting",
-    )
-    scheduler = get_scheduler_from_request(InputData, "inpaint")
-    (lora_weights, lora_hf_id) = get_lora_params(frozen_args.use_lora)
-
-    init_image = decode_base64_to_image(InputData.image)
-    mask = decode_base64_to_image(InputData.mask)
-
-    print(
-        f"Prompt: {InputData.prompt}, "
-        f'Negative Prompt: {InputData.negative_prompt}", '
-        f'Seed: {InputData.seed}", '
-        f"Model: {model_id}, "
-        f"Scheduler: {scheduler}."
-    )
-
-    res = inpaint_inf(
-        InputData.prompt,
-        InputData.negative_prompt,
-        {"image": init_image, "mask": mask},
-        InputData.height,
-        InputData.width,
-        InputData.is_full_res,
-        InputData.full_res_padding,
-        InputData.steps,
-        InputData.cfg_scale,
-        InputData.seed,
-        batch_count=InputData.n_iter,
-        batch_size=1,
-        scheduler=scheduler,
-        model_id=model_id,
-        custom_vae=frozen_args.custom_vae or "None",
-        precision="fp16",
-        device=get_device(frozen_args.device),
-        max_length=frozen_args.max_length,
-        save_metadata_to_json=frozen_args.save_metadata_to_json,
-        save_metadata_to_png=frozen_args.write_metadata_to_png,
-        lora_weights=lora_weights,
-        lora_hf_id=lora_hf_id,
-        ondemand=frozen_args.ondemand,
-        repeatable_seeds=False,
-    )
-
-    # Since we're not streaming we just want the last generator result
-    for items_so_far in res:
-        items = items_so_far
-
-    return {
-        "images": encode_pil_to_base64(items[0]),
-        "parameters": {},
-        "info": items[1],
-    }
-
-
-# Rest API: /sdapi/v1/outpaint (Outpainting)
-class DirectionParam(str, Enum):
-    left = "left"
-    right = "right"
-    up = "up"
-    down = "down"
-
-
-class OutpaintInputData(GenerationInputData):
-    init_images: list[str]
-    pixels: int = Field(
-        default=frozen_args.pixels, ge=8, le=256, multiple_of=8
-    )
-    mask_blur: int = Field(default=frozen_args.mask_blur, ge=0, le=64)
-    directions: set[DirectionParam] = [
-        direction
-        for direction in ["left", "right", "up", "down"]
-        if vars(frozen_args)[direction]
-    ]
-    noise_q: float = frozen_args.noise_q
-    color_variation: float = frozen_args.color_variation
-    override_settings: PaintModelOverideSettings = None
-
-
-@sdapi.post(
-    "/v1/outpaint",
-    summary="Does outpainting generation on an image",
-    response_model=GenerationResponseData,
-)
-def outpaint_api(
-    InputData: OutpaintInputData,
-):
-    model_id = get_model_from_request(
-        InputData,
-        checkpoint_type="inpainting",
-        fallback_model="stabilityai/stable-diffusion-2-inpainting",
-    )
-    scheduler = get_scheduler_from_request(InputData, "outpaint")
-    (lora_weights, lora_hf_id) = get_lora_params(frozen_args.use_lora)
-
-    init_image = decode_base64_to_image(InputData.init_images[0])
-
-    print(
-        f"Prompt: {InputData.prompt}, "
-        f"Negative Prompt: {InputData.negative_prompt}, "
-        f"Seed: {InputData.seed}, "
-        f"Model: {model_id}, "
-        f"Scheduler: {scheduler}."
-    )
-
-    res = outpaint_inf(
-        InputData.prompt,
-        InputData.negative_prompt,
-        init_image,
-        InputData.pixels,
-        InputData.mask_blur,
-        InputData.directions,
-        InputData.noise_q,
-        InputData.color_variation,
-        InputData.height,
-        InputData.width,
-        InputData.steps,
-        InputData.cfg_scale,
-        InputData.seed,
-        batch_count=InputData.n_iter,
-        batch_size=1,
-        scheduler=scheduler,
-        model_id=model_id,
-        custom_vae=frozen_args.custom_vae or "None",
-        precision="fp16",
-        device=get_device(frozen_args.device),
-        max_length=frozen_args.max_length,
-        save_metadata_to_json=frozen_args.save_metadata_to_json,
-        save_metadata_to_png=frozen_args.write_metadata_to_png,
-        lora_weights=lora_weights,
-        lora_hf_id=lora_hf_id,
-        ondemand=frozen_args.ondemand,
-        repeatable_seeds=False,
-    )
-
-    # Since we're not streaming we just want the last generator result
-    for items_so_far in res:
-        items = items_so_far
-
-    return {
-        "images": encode_pil_to_base64(items[0]),
-        "parameters": {},
-        "info": items[1],
-    }
-
-
-# Rest API: /sdapi/v1/upscaler (Upscaling)
-class UpscalerModelOverideSettings(BaseModel):
-    sd_model_checkpoint: str = get_model_from_request(
-        checkpoint_type="upscaler",
-        fallback_model="stabilityai/stable-diffusion-x4-upscaler",
-    )
-
-
-class UpscalerInputData(GenerationInputData):
-    init_images: list[str] = Field(
-        description="Base64 encoded image to upscale"
-    )
-    noise_level: int = frozen_args.noise_level
-    override_settings: UpscalerModelOverideSettings = None
-
-
-@sdapi.post(
-    "/v1/upscaler",
-    summary="Does image upscaling",
-    response_model=GenerationResponseData,
-)
-def upscaler_api(
-    InputData: UpscalerInputData,
-):
-    model_id = get_model_from_request(
-        InputData,
-        checkpoint_type="upscaler",
-        fallback_model="stabilityai/stable-diffusion-x4-upscaler",
-    )
-    scheduler = get_scheduler_from_request(InputData, "upscaler")
-    (lora_weights, lora_hf_id) = get_lora_params(frozen_args.use_lora)
-
-    init_image = decode_base64_to_image(InputData.init_images[0])
-
-    print(
-        f"Prompt: {InputData.prompt}, "
-        f"Negative Prompt: {InputData.negative_prompt}, "
-        f"Seed: {InputData.seed}, "
-        f"Model: {model_id}, "
-        f"Scheduler: {scheduler}."
-    )
-
-    res = upscaler_inf(
-        InputData.prompt,
-        InputData.negative_prompt,
-        init_image,
-        InputData.height,
-        InputData.width,
-        InputData.steps,
-        InputData.noise_level,
-        InputData.cfg_scale,
-        InputData.seed,
-        batch_count=InputData.n_iter,
-        batch_size=1,
-        scheduler=scheduler,
-        model_id=model_id,
-        custom_vae=frozen_args.custom_vae or "None",
-        precision="fp16",
-        device=get_device(frozen_args.device),
-        max_length=frozen_args.max_length,
-        save_metadata_to_json=frozen_args.save_metadata_to_json,
-        save_metadata_to_png=frozen_args.write_metadata_to_png,
-        lora_weights=lora_weights,
-        lora_hf_id=lora_hf_id,
-        ondemand=frozen_args.ondemand,
-        repeatable_seeds=False,
-    )
-
-    # Since we're not streaming we just want the last generator result
-    for items_so_far in res:
-        items = items_so_far
-
-    return {
-        "images": encode_pil_to_base64(items[0]),
-        "parameters": {},
-        "info": items[1],
-    }