Refactor SHARK Studio on top of SHARK-Turbine

monorimet commented 10 months ago

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

dan-garvey commented 10 months ago

gpetters94 commented 10 months ago

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

one-lithe-rune commented 10 months ago

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.

monorimet commented 10 months ago

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

nod-ai / SHARK