SHARK icon indicating copy to clipboard operation
SHARK copied to clipboard

Refactor SHARK Studio on top of SHARK-Turbine

Open monorimet opened this issue 1 year ago • 4 comments

Motivation

As the SHARK Studio application and its scope grows, we notice an opportunity to rework the general code structure of the application from a stable diffusion webui to a multi-modal generative AI studio.

Goals

  • refactor the app (and, generally, the repository) on top of SHARK-Turbine for model import and compilation
  • unified APIs for each Studio feature
  • shift code structure to more intuitive design, encouraging contribution and code health.
  • enable more robust CI coverage (testing, deployment)
  • minimize statefulness and global variables...

Outline

Currently the code for the UI / app spec is structured as follows:

  • SHARK
    • apps/
      • stable_diffusion/
        • shark_sd.spec <-- used by PyInstaller in our nightly build via pyinstaller shark_sd.spec to build shark_studio.exe
        • shark_sd_cli.spec <-- useful for users/devs who want a terminal-only shark_studio.exe
        • studio_bundle.spec <-- useful for debugging discrepancies between python venv invocations and pyinstaller executables.
        • shark_studio_imports.py <-- controls the packaging of the studio application and which pip package datas+metadatas are shipped with the studio.exe
        • web/ <--- this entire directory should be app-level, not model or mode-level!
          • ui/
          • utils/
          • index.py <-- main entry point for users, the studio.exe essentially wraps this python file and all its dependencies.
        • src/
          • models/
          • utils/
          • pipelines/
          • schedulers/
      • language_models/ <-- contains model wrappers and inference scripts that plug into the chatbot UI
        • shark_llama_cli.spec
        • scripts/
        • src/
          • model_wrappers/
          • pipelines/
        • utils.py

The proposed new layout is outlined as follows:

  • SHARK
    • apps/
      • shark_studio/
        • src/
          • fetch compiled modules from turbine
            • takes nn.modules, HF IDs, .mlir and gives compiled module(s)
          • Abstracted wrapper for modular cases (hold my compiled modules please?) (shark_inference 2.0)
          • api (shared runtime api -- webui + cli + rest API)
            • stable_diffusion API (match auto1111)
            • chatbot
            • other studio features
        • app/
          • index.py
          • ui/ <-- gradio code for each tab, only calls to functions and methods from src/
        • pyinstaller .specs

Tasks

  • [ ] Move pyinstaller specs + UI to SHARK/apps/shark_studio/ (@monorimet )
  • [ ] Add test framework for shark studio (@monorimet )
  • [x] API -- move to SHARK/apps/shark_studio/src/api/, one file for each Studio functionality (each tab is a good start)
    • [x] API - stable diffusion (must enable reuse of existing auto1111 REST API plugin support! )
    • [x] API - chatbot
  • [ ] Write a test suite on the unified API for existing use cases of SHARK Studio (@monorimet )
  • [x] Turbine integration (Chatbot) -- the fabled stateless llama (@dan-garvey )
  • [ ] Turbine integration (Stable Diffusion)
  • [ ] More turbine integration: everything but the runtime

monorimet avatar Nov 01 '23 21:11 monorimet

image

dan-garvey avatar Nov 01 '23 21:11 dan-garvey

LGTM. I wonder if we should rename this repo to SHARK_studio and break any future apps into separate repos.

gpetters94 avatar Nov 02 '23 01:11 gpetters94

I think it needs docs folders.

Probably for the individual interfaces under api (i.e the connecting to blender one, and the one I haven't finished yet for koboldcpp), and for ui (current webui really needs documentation, future ui should have some).

I don't understand the pipeline guts enough to know where any docs for src in general should go.

one-lithe-rune avatar Nov 02 '23 09:11 one-lithe-rune

https://github.com/nod-ai/SHARK-Turbine/issues/152

monorimet avatar Nov 03 '23 17:11 monorimet