Fluminato

.fluminaignore

@@ -0,0 +1,166 @@
+.git/
+.gitmodules
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  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/
+.DS_Store

.gitattributes

@@ -0,0 +1,5 @@
+flux1-dev.safetensors filter=lfs diff=lfs merge=lfs -text
+ae.safetensors filter=lfs diff=lfs merge=lfs -text
+t5-v1_1-xxl-encoder-bf16//*.safetensor filter=lfs diff=lfs merge=lfs -text
+t5-v1_1-xxl-encoder-bf16//spiece.model filter=lfs diff=lfs merge=lfs -text
+t5-v1_1-xxl-encoder-bf16/model.safetensors filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,20 @@
+{
+  "config_path": "configs/config-dev.json",
+  "flow_model_path": "flux1-dev.safetensors",
+  "model_version": "flux-dev",
+  "flux_device": "cuda",
+  "autoencoder_path": null,
+  "autoencoder_device": "cuda",
+  "text_enc_path": null,
+  "text_enc_device": "cuda",
+  "num_to_quant": 20,
+  "compile": false,
+  "quant_text_enc": "qfloat8",
+  "quant_ae": false,
+  "offload_flow": false,
+  "offload_ae": true,
+  "offload_text_enc": true,
+  "prequantized_flow": false,
+  "quantize_modulation": true,
+  "quantize_flow_embedder_layers": false
+}

configs/config-dev-1-RTX6000ADA.json

@@ -1,57 +0,0 @@
-{
-  "version": "flux-dev",
-  "params": {
-    "in_channels": 64,
-    "vec_in_dim": 768,
-    "context_in_dim": 4096,
-    "hidden_size": 3072,
-    "mlp_ratio": 4.0,
-    "num_heads": 24,
-    "depth": 19,
-    "depth_single_blocks": 38,
-    "axes_dim": [
-      16,
-      56,
-      56
-    ],
-    "theta": 10000,
-    "qkv_bias": true,
-    "guidance_embed": true
-  },
-  "ae_params": {
-    "resolution": 256,
-    "in_channels": 3,
-    "ch": 128,
-    "out_ch": 3,
-    "ch_mult": [
-      1,
-      2,
-      4,
-      4
-    ],
-    "num_res_blocks": 2,
-    "z_channels": 16,
-    "scale_factor": 0.3611,
-    "shift_factor": 0.1159
-  },
-  "ckpt_path": "/big/generator-ui/flux-testing/flux/model-dir/flux1-dev.sft",
-  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir/ae.sft",
-  "repo_id": "black-forest-labs/FLUX.1-dev",
-  "repo_flow": "flux1-dev.sft",
-  "repo_ae": "ae.sft",
-  "text_enc_max_length": 512,
-  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
-  "text_enc_device": "cuda:0",
-  "ae_device": "cuda:0",
-  "flux_device": "cuda:0",
-  "flow_dtype": "float16",
-  "ae_dtype": "bfloat16",
-  "text_enc_dtype": "bfloat16",
-  "flow_quantization_dtype": "qfloat8",
-  "text_enc_quantization_dtype": "qfloat8",
-  "compile_extras": true,
-  "compile_blocks": true,
-  "offload_text_encoder": false,
-  "offload_vae": false,
-  "offload_flow": false
-}

configs/config-dev-cuda0.json

@@ -1,56 +0,0 @@
-{
-  "version": "flux-dev",
-  "params": {
-    "in_channels": 64,
-    "vec_in_dim": 768,
-    "context_in_dim": 4096,
-    "hidden_size": 3072,
-    "mlp_ratio": 4.0,
-    "num_heads": 24,
-    "depth": 19,
-    "depth_single_blocks": 38,
-    "axes_dim": [
-      16,
-      56,
-      56
-    ],
-    "theta": 10000,
-    "qkv_bias": true,
-    "guidance_embed": true
-  },
-  "ae_params": {
-    "resolution": 256,
-    "in_channels": 3,
-    "ch": 128,
-    "out_ch": 3,
-    "ch_mult": [
-      1,
-      2,
-      4,
-      4
-    ],
-    "num_res_blocks": 2,
-    "z_channels": 16,
-    "scale_factor": 0.3611,
-    "shift_factor": 0.1159
-  },
-  "ckpt_path": "/big/generator-ui/flux-testing/flux/model-dir/flux1-dev.sft",
-  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir/ae.sft",
-  "repo_id": "black-forest-labs/FLUX.1-dev",
-  "repo_flow": "flux1-dev.sft",
-  "repo_ae": "ae.sft",
-  "text_enc_max_length": 512,
-  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
-  "text_enc_device": "cuda:0",
-  "ae_device": "cuda:0",
-  "flux_device": "cuda:0",
-  "flow_dtype": "float16",
-  "ae_dtype": "bfloat16",
-  "text_enc_dtype": "bfloat16",
-  "text_enc_quantization_dtype": "qfloat8",
-  "compile_extras": false,
-  "compile_blocks": false,
-  "offload_ae": false,
-  "offload_text_enc": false,
-  "offload_flow": false
-}

configs/config-dev-eval.json

@@ -1,57 +0,0 @@
-{
-  "version": "flux-dev",
-  "params": {
-    "in_channels": 64,
-    "vec_in_dim": 768,
-    "context_in_dim": 4096,
-    "hidden_size": 3072,
-    "mlp_ratio": 4.0,
-    "num_heads": 24,
-    "depth": 19,
-    "depth_single_blocks": 38,
-    "axes_dim": [
-      16,
-      56,
-      56
-    ],
-    "theta": 10000,
-    "qkv_bias": true,
-    "guidance_embed": true
-  },
-  "ae_params": {
-    "resolution": 256,
-    "in_channels": 3,
-    "ch": 128,
-    "out_ch": 3,
-    "ch_mult": [
-      1,
-      2,
-      4,
-      4
-    ],
-    "num_res_blocks": 2,
-    "z_channels": 16,
-    "scale_factor": 0.3611,
-    "shift_factor": 0.1159
-  },
-  "ckpt_path": "/big/generator-ui/flux-testing/flux/model-dir/flux1-dev.sft",
-  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir/ae.sft",
-  "repo_id": "black-forest-labs/FLUX.1-dev",
-  "repo_flow": "flux1-dev.sft",
-  "repo_ae": "ae.sft",
-  "text_enc_max_length": 512,
-  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
-  "text_enc_device": "cuda:1",
-  "ae_device": "cuda:1",
-  "flux_device": "cuda:0",
-  "flow_dtype": "float16",
-  "ae_dtype": "bfloat16",
-  "text_enc_dtype": "bfloat16",
-  "flow_quantization_dtype": "qfloat8",
-  "text_enc_quantization_dtype": "qfloat8",
-  "compile_extras": false,
-  "compile_blocks": false,
-  "offload_ae": false,
-  "offload_text_enc": false,
-  "offload_flow": false
-}

configs/config-dev-gigaquant.json

@@ -1,58 +0,0 @@
-{
-  "version": "flux-dev",
-  "params": {
-    "in_channels": 64,
-    "vec_in_dim": 768,
-    "context_in_dim": 4096,
-    "hidden_size": 3072,
-    "mlp_ratio": 4.0,
-    "num_heads": 24,
-    "depth": 19,
-    "depth_single_blocks": 38,
-    "axes_dim": [
-      16,
-      56,
-      56
-    ],
-    "theta": 10000,
-    "qkv_bias": true,
-    "guidance_embed": true
-  },
-  "ae_params": {
-    "resolution": 256,
-    "in_channels": 3,
-    "ch": 128,
-    "out_ch": 3,
-    "ch_mult": [
-      1,
-      2,
-      4,
-      4
-    ],
-    "num_res_blocks": 2,
-    "z_channels": 16,
-    "scale_factor": 0.3611,
-    "shift_factor": 0.1159
-  },
-  "ckpt_path": "/big/generator-ui/flux-testing/flux/model-dir/flux1-dev.sft",
-  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir/ae.sft",
-  "repo_id": "black-forest-labs/FLUX.1-dev",
-  "repo_flow": "flux1-dev.sft",
-  "repo_ae": "ae.sft",
-  "text_enc_max_length": 512,
-  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
-  "text_enc_device": "cuda:0",
-  "ae_device": "cuda:0",
-  "flux_device": "cuda:0",
-  "flow_dtype": "float16",
-  "ae_dtype": "bfloat16",
-  "text_enc_dtype": "bfloat16",
-  "num_to_quant": 220,
-  "flow_quantization_dtype": "qint4",
-  "text_enc_quantization_dtype": "qint4",
-  "ae_quantization_dtype": "qint4",
-  "clip_quantization_dtype": "qint4",
-  "compile_extras": false,
-  "compile_blocks": false,
-  "quantize_extras": true
-}

configs/config-dev-offload-1-4080.json

@@ -1,58 +0,0 @@
-{
-  "version": "flux-dev",
-  "params": {
-    "in_channels": 64,
-    "vec_in_dim": 768,
-    "context_in_dim": 4096,
-    "hidden_size": 3072,
-    "mlp_ratio": 4.0,
-    "num_heads": 24,
-    "depth": 19,
-    "depth_single_blocks": 38,
-    "axes_dim": [
-      16,
-      56,
-      56
-    ],
-    "theta": 10000,
-    "qkv_bias": true,
-    "guidance_embed": true
-  },
-  "ae_params": {
-    "resolution": 256,
-    "in_channels": 3,
-    "ch": 128,
-    "out_ch": 3,
-    "ch_mult": [
-      1,
-      2,
-      4,
-      4
-    ],
-    "num_res_blocks": 2,
-    "z_channels": 16,
-    "scale_factor": 0.3611,
-    "shift_factor": 0.1159
-  },
-  "ckpt_path": "/big/generator-ui/flux-testing/flux/model-dir/flux1-dev.sft",
-  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir/ae.sft",
-  "repo_id": "black-forest-labs/FLUX.1-dev",
-  "repo_flow": "flux1-dev.sft",
-  "repo_ae": "ae.sft",
-  "text_enc_max_length": 512,
-  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
-  "text_enc_device": "cuda:0",
-  "ae_device": "cuda:0",
-  "flux_device": "cuda:0",
-  "flow_dtype": "float16",
-  "ae_dtype": "bfloat16",
-  "text_enc_dtype": "bfloat16",
-  "flow_quantization_dtype": "qfloat8",
-  "text_enc_quantization_dtype": "qint4",
-  "ae_quantization_dtype": "qfloat8",
-  "compile_extras": true,
-  "compile_blocks": true,
-  "offload_text_encoder": true,
-  "offload_vae": true,
-  "offload_flow": true
-}

configs/config-dev-offload-1-4090.json

@@ -1,58 +0,0 @@
-{
-  "version": "flux-dev",
-  "params": {
-    "in_channels": 64,
-    "vec_in_dim": 768,
-    "context_in_dim": 4096,
-    "hidden_size": 3072,
-    "mlp_ratio": 4.0,
-    "num_heads": 24,
-    "depth": 19,
-    "depth_single_blocks": 38,
-    "axes_dim": [
-      16,
-      56,
-      56
-    ],
-    "theta": 10000,
-    "qkv_bias": true,
-    "guidance_embed": true
-  },
-  "ae_params": {
-    "resolution": 256,
-    "in_channels": 3,
-    "ch": 128,
-    "out_ch": 3,
-    "ch_mult": [
-      1,
-      2,
-      4,
-      4
-    ],
-    "num_res_blocks": 2,
-    "z_channels": 16,
-    "scale_factor": 0.3611,
-    "shift_factor": 0.1159
-  },
-  "ckpt_path": "/big/generator-ui/flux-testing/flux/model-dir/flux1-dev.sft",
-  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir/ae.sft",
-  "repo_id": "black-forest-labs/FLUX.1-dev",
-  "repo_flow": "flux1-dev.sft",
-  "repo_ae": "ae.sft",
-  "text_enc_max_length": 512,
-  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
-  "text_enc_device": "cuda:0",
-  "ae_device": "cuda:0",
-  "flux_device": "cuda:0",
-  "flow_dtype": "float16",
-  "ae_dtype": "bfloat16",
-  "text_enc_dtype": "bfloat16",
-  "flow_quantization_dtype": "qfloat8",
-  "text_enc_quantization_dtype": "qint4",
-  "ae_quantization_dtype": "qfloat8",
-  "compile_extras": true,
-  "compile_blocks": true,
-  "offload_text_encoder": true,
-  "offload_vae": true,
-  "offload_flow": false
-}

configs/config-dev-offload.json

@@ -1,58 +0,0 @@
-{
-  "version": "flux-dev",
-  "params": {
-    "in_channels": 64,
-    "vec_in_dim": 768,
-    "context_in_dim": 4096,
-    "hidden_size": 3072,
-    "mlp_ratio": 4.0,
-    "num_heads": 24,
-    "depth": 19,
-    "depth_single_blocks": 38,
-    "axes_dim": [
-      16,
-      56,
-      56
-    ],
-    "theta": 10000,
-    "qkv_bias": true,
-    "guidance_embed": true
-  },
-  "ae_params": {
-    "resolution": 256,
-    "in_channels": 3,
-    "ch": 128,
-    "out_ch": 3,
-    "ch_mult": [
-      1,
-      2,
-      4,
-      4
-    ],
-    "num_res_blocks": 2,
-    "z_channels": 16,
-    "scale_factor": 0.3611,
-    "shift_factor": 0.1159
-  },
-  "ckpt_path": "/big/generator-ui/flux-testing/flux/model-dir/flux1-dev.sft",
-  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir/ae.sft",
-  "repo_id": "black-forest-labs/FLUX.1-dev",
-  "repo_flow": "flux1-dev.sft",
-  "repo_ae": "ae.sft",
-  "text_enc_max_length": 512,
-  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
-  "text_enc_device": "cuda:0",
-  "ae_device": "cuda:0",
-  "flux_device": "cuda:0",
-  "flow_dtype": "float16",
-  "ae_dtype": "bfloat16",
-  "text_enc_dtype": "bfloat16",
-  "flow_quantization_dtype": "qfloat8",
-  "text_enc_quantization_dtype": "qint4",
-  "ae_quantization_dtype": "qfloat8",
-  "compile_extras": false,
-  "compile_blocks": false,
-  "offload_text_encoder": true,
-  "offload_vae": true,
-  "offload_flow": true
-}

configs/config-dev-prequant.json

@@ -1,57 +0,0 @@
-{
-  "version": "flux-dev",
-  "params": {
-    "in_channels": 64,
-    "vec_in_dim": 768,
-    "context_in_dim": 4096,
-    "hidden_size": 3072,
-    "mlp_ratio": 4.0,
-    "num_heads": 24,
-    "depth": 19,
-    "depth_single_blocks": 38,
-    "axes_dim": [
-      16,
-      56,
-      56
-    ],
-    "theta": 10000,
-    "qkv_bias": true,
-    "guidance_embed": true
-  },
-  "ae_params": {
-    "resolution": 256,
-    "in_channels": 3,
-    "ch": 128,
-    "out_ch": 3,
-    "ch_mult": [
-      1,
-      2,
-      4,
-      4
-    ],
-    "num_res_blocks": 2,
-    "z_channels": 16,
-    "scale_factor": 0.3611,
-    "shift_factor": 0.1159
-  },
-  "ckpt_path": "/big/generator-ui/flux-testing/flux/flux-fp16-acc/flux_fp8.safetensors",
-  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir/ae.sft",
-  "repo_id": "black-forest-labs/FLUX.1-dev",
-  "repo_flow": "flux1-dev.sft",
-  "repo_ae": "ae.sft",
-  "text_enc_max_length": 512,
-  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
-  "text_enc_device": "cuda:1",
-  "ae_device": "cuda:1",
-  "flux_device": "cuda:0",
-  "flow_dtype": "float16",
-  "ae_dtype": "bfloat16",
-  "text_enc_dtype": "bfloat16",
-  "text_enc_quantization_dtype": "qfloat8",
-  "compile_extras": false,
-  "compile_blocks": false,
-  "prequantized_flow": true,
-  "offload_ae": false,
-  "offload_text_enc": false,
-  "offload_flow": false
-}

configs/config-dev.json

@@ -34,16 +34,16 @@
     "scale_factor": 0.3611,
     "shift_factor": 0.1159
   },
-  "ckpt_path": "/big/generator-ui/flux-testing/flux/model-dir/flux1-dev.sft",
-  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir/ae.sft",
+  "ckpt_path": "flux1-dev.safetensors",
+  "ae_path": "ae.safetensors",
   "repo_id": "black-forest-labs/FLUX.1-dev",
-  "repo_flow": "flux1-dev.sft",
-  "repo_ae": "ae.sft",
+  "repo_flow": "flux1-dev.safetensors",
+  "repo_ae": "ae.safetensors",
   "text_enc_max_length": 512,
-  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
-  "text_enc_device": "cuda:1",
-  "ae_device": "cuda:1",
-  "flux_device": "cuda:0",
+  "text_enc_path": "t5-v1_1-xxl-encoder-bf16",
+  "text_enc_device": "cuda",
+  "ae_device": "cuda",
+  "flux_device": "cuda",
   "flow_dtype": "float16",
   "ae_dtype": "bfloat16",
   "text_enc_dtype": "bfloat16",
@@ -54,4 +54,4 @@
   "offload_ae": false,
   "offload_text_enc": false,
   "offload_flow": false
-}
+}

configs/config-schnell-cuda0.json

@@ -1,57 +0,0 @@
-{
-  "version": "flux-schnell",
-  "params": {
-    "in_channels": 64,
-    "vec_in_dim": 768,
-    "context_in_dim": 4096,
-    "hidden_size": 3072,
-    "mlp_ratio": 4.0,
-    "num_heads": 24,
-    "depth": 19,
-    "depth_single_blocks": 38,
-    "axes_dim": [
-      16,
-      56,
-      56
-    ],
-    "theta": 10000,
-    "qkv_bias": true,
-    "guidance_embed": false
-  },
-  "ae_params": {
-    "resolution": 256,
-    "in_channels": 3,
-    "ch": 128,
-    "out_ch": 3,
-    "ch_mult": [
-      1,
-      2,
-      4,
-      4
-    ],
-    "num_res_blocks": 2,
-    "z_channels": 16,
-    "scale_factor": 0.3611,
-    "shift_factor": 0.1159
-  },
-  "ckpt_path": "/big/generator-ui/flux-testing/flux/model-dir-schnell/flux1-schnell.sft",
-  "ae_path": "/big/generator-ui/flux-testing/flux/model-dir-schnell/ae.sft",
-  "repo_id": "black-forest-labs/FLUX.1-schnell",
-  "repo_flow": "flux1-schnell.sft",
-  "repo_ae": "ae.sft",
-  "text_enc_max_length": 256,
-  "text_enc_path": "city96/t5-v1_1-xxl-encoder-bf16",
-  "text_enc_device": "cuda:0",
-  "ae_device": "cuda:0",
-  "flux_device": "cuda:0",
-  "flow_dtype": "float16",
-  "ae_dtype": "bfloat16",
-  "text_enc_dtype": "bfloat16",
-  "text_enc_quantization_dtype": "qfloat8",
-  "ae_quantization_dtype": "qfloat8",
-  "compile_extras": false,
-  "compile_blocks": false,
-  "offload_ae": false,
-  "offload_text_enc": false,
-  "offload_flow": false
-}

fireworks.json

@@ -0,0 +1 @@
+{"_is_flumina_model": true}

flumina.py

@@ -0,0 +1,326 @@
+# flumina.py
+import torch
+import io
+import json
+from fireworks.flumina import FluminaModule, main as flumina_main
+from fireworks.flumina.route import post
+import pydantic
+from pydantic import BaseModel
+from fastapi import Header
+from fastapi.responses import Response
+import math
+import os
+import re
+import PIL.Image as Image
+from typing import Optional, Set, Tuple
+
+from flux_pipeline import FluxPipeline
+from util import load_config, ModelVersion
+
+# Util
+def _aspect_ratio_to_width_height(aspect_ratio: str) -> Tuple[int, int]:
+    """
+    Convert specified aspect ratio to a height/width pair.
+    """
+    if ":" not in aspect_ratio:
+        raise ValueError(
+            f"Invalid aspect ratio: {aspect_ratio}. Aspect ratio must be in w:h format, e.g. 16:9"
+        )
+
+    w, h = aspect_ratio.split(":")
+    try:
+        w, h = int(w), int(h)
+    except ValueError:
+        raise ValueError(
+            f"Invalid aspect ratio: {aspect_ratio}. Aspect ratio must be in w:h format, e.g. 16:9"
+        )
+
+    valid_aspect_ratios = [
+        (1, 1),
+        (21, 9),
+        (16, 9),
+        (3, 2),
+        (5, 4),
+        (4, 5),
+        (2, 3),
+        (9, 16),
+        (9, 21),
+    ]
+    if (w, h) not in valid_aspect_ratios:
+        raise ValueError(
+            f"Invalid aspect ratio: {aspect_ratio}. Aspect ratio must be one of {valid_aspect_ratios}"
+        )
+
+    # We consider megapixel not 10^6 pixels but 2^20 (1024x1024) pixels
+    TARGET_SIZE_MP = 1
+    target_size = TARGET_SIZE_MP * 2**20
+
+    width = math.sqrt(target_size / (w * h)) * w
+    height = math.sqrt(target_size / (w * h)) * h
+
+    PAD_MULTIPLE = 64
+
+    if PAD_MULTIPLE:
+        width = width // PAD_MULTIPLE * PAD_MULTIPLE
+        height = height // PAD_MULTIPLE * PAD_MULTIPLE
+
+    return int(width), int(height)
+
+
+def encode_image(
+    image: Image.Image, mime_type: str, jpeg_quality: int = 95
+) -> bytes:
+    buffered = io.BytesIO()
+    if mime_type == "image/jpeg":
+        if jpeg_quality < 0 or jpeg_quality > 100:
+            raise ValueError(
+                f"jpeg_quality must be between 0 and 100, not {jpeg_quality}"
+            )
+        image.save(buffered, format="JPEG", quality=jpeg_quality)
+    elif mime_type == "image/png":
+        image.save(buffered, format="PNG")
+    else:
+        raise ValueError(f"invalid mime_type {mime_type}")
+    return buffered.getvalue()
+
+
+def parse_accept_header(accept: str) -> str:
+    # Split the string into the comma-separated components
+    parts = accept.split(",")
+    weighted_types = []
+
+    for part in parts:
+        # Use a regular expression to extract the media type and the optional q-factor
+        match = re.match(
+            r"(?P<media_type>[^;]+)(;q=(?P<q_factor>\d+(\.\d+)?))?", part.strip()
+        )
+        if match:
+            media_type = match.group("media_type")
+            q_factor = (
+                float(match.group("q_factor")) if match.group("q_factor") else 1.0
+            )
+            weighted_types.append((media_type, q_factor))
+        else:
+            raise ValueError(f"Malformed Accept header value: {part.strip()}")
+
+    # Sort the media types by q-factor, descending
+    sorted_types = sorted(weighted_types, key=lambda x: x[1], reverse=True)
+
+    for media_type, _ in sorted_types:
+        if media_type in {"image/jpeg", "image/png"}:
+            return media_type
+
+    raise ValueError(f"Accept header did not include ones of supported MIME types: image/jpeg, image/png")
+
+
+# Define request and response schemata
+class Text2ImageRequest(BaseModel):
+    prompt: str
+    aspect_ratio: str = "16:9"
+    guidance_scale: float = 3.5
+    num_inference_steps: int = 30
+    seed: int = 0
+
+
+class Error(BaseModel):
+    object: str = "error"
+    type: str = "invalid_request_error"
+    message: str
+
+
+class ErrorResponse(BaseModel):
+    error: Error = pydantic.Field(default_factory=Error)
+
+
+class BillingInfo(BaseModel):
+    steps: int
+    height: int
+    width: int
+    is_control_net: bool = False
+
+
+class FluminaModule(FluminaModule):
+    def __init__(self):
+        super().__init__()
+
+        # Read configuration from config.json
+        with open('config.json', 'r') as f:
+            config_data = json.load(f)
+
+        # Now, we need to construct the config and load the model
+        if 'config_path' in config_data:
+            self.pipeline = FluxPipeline.load_pipeline_from_config_path(
+                config_data['config_path'],
+                flow_model_path=config_data.get('flow_model_path', None)
+            )
+        else:
+            model_version = (
+                ModelVersion.flux_dev
+                if config_data.get('model_version', 'flux-dev') == "flux-dev"
+                else ModelVersion.flux_schnell
+            )
+            config = load_config(
+                model_version,
+                flux_path=config_data.get('flow_model_path', None),
+                flux_device=config_data.get('flux_device', 'cuda:0'),
+                ae_path=config_data.get('autoencoder_path', None),
+                ae_device=config_data.get('autoencoder_device', 'cuda:0'),
+                text_enc_path=config_data.get('text_enc_path', None),
+                text_enc_device=config_data.get('text_enc_device', 'cuda:0'),
+                flow_dtype="float16",
+                text_enc_dtype="bfloat16",
+                ae_dtype="bfloat16",
+                num_to_quant=config_data.get('num_to_quant', 20),
+                compile_extras=config_data.get('compile', False),
+                compile_blocks=config_data.get('compile', False),
+                quant_text_enc=(
+                    None
+                    if config_data.get('quant_text_enc', 'qfloat8') == "bf16"
+                    else config_data.get('quant_text_enc', 'qfloat8')
+                ),
+                quant_ae=config_data.get('quant_ae', False),
+                offload_flow=config_data.get('offload_flow', False),
+                offload_ae=config_data.get('offload_ae', True),
+                offload_text_enc=config_data.get('offload_text_enc', True),
+                prequantized_flow=config_data.get('prequantized_flow', False),
+                quantize_modulation=config_data.get('quantize_modulation', True),
+                quantize_flow_embedder_layers=config_data.get(
+                    'quantize_flow_embedder_layers', False
+                ),
+            )
+            self.pipeline = FluxPipeline.load_pipeline_from_config(config)
+
+        # Initialize LoRA adapters
+        self.lora_adapters: Set[str] = set()
+        self.active_lora_adapter: Optional[str] = None
+        self._test_return_sync_response = False
+
+    def _error_response(self, code: int, message: str) -> Response:
+        response_json = ErrorResponse(
+            error=Error(message=message),
+        ).json()
+        if self._test_return_sync_response:
+            return response_json
+        else:
+            return Response(
+                response_json,
+                status_code=code,
+                media_type="application/json",
+            )
+
+    def _image_response(
+        self, image_bytes: bytes, mime_type: str, billing_info: BillingInfo
+    ):
+        if self._test_return_sync_response:
+            return image_bytes
+        else:
+            headers = {'Fireworks-Billing-Properties': billing_info.json()}
+            return Response(
+                image_bytes, status_code=200, media_type=mime_type, headers=headers
+            )
+
+    @post('/text_to_image')
+    async def text_to_image(
+        self,
+        body: Text2ImageRequest,
+        accept: str = Header("image/jpeg"),
+    ):
+        mime_type = parse_accept_header(accept)
+        width, height = _aspect_ratio_to_width_height(body.aspect_ratio)
+        img_bio = self.pipeline.generate(
+            prompt=body.prompt,
+            height=height,
+            width=width,
+            guidance=body.guidance_scale,
+            num_steps=body.num_inference_steps,
+            seed=body.seed,
+        )
+
+        billing_info = BillingInfo(
+            steps=body.num_inference_steps,
+            height=height,
+            width=width,
+        )
+        return self._image_response(img_bio.getvalue(), mime_type, billing_info)
+
+    @property
+    def supported_addon_types(self):
+        return ['lora']
+
+    # Addon interface methods adjusted to remove ControlNet support
+    def load_addon(
+        self,
+        addon_account_id: str,
+        addon_model_id: str,
+        addon_type: str,
+        addon_data_path: os.PathLike,
+    ):
+        if addon_type not in self.supported_addon_types:
+            raise ValueError(
+                f"Invalid addon type {addon_type}. Supported types: {self.supported_addon_types}"
+            )
+
+        qualname = f"accounts/{addon_account_id}/models/{addon_model_id}"
+
+        if addon_type == 'lora':
+            self.pipeline.load_lora_weights(addon_data_path, adapter_name=qualname)
+            self.lora_adapters.add(qualname)
+        else:
+            raise NotImplementedError(
+                f'Addon support for type {addon_type} not implemented'
+            )
+
+    def unload_addon(
+        self, addon_account_id: str, addon_model_id: str, addon_type: str
+    ):
+        qualname = f"accounts/{addon_account_id}/models/{addon_model_id}"
+
+        if addon_type == 'lora':
+            assert qualname in self.lora_adapters
+            self.pipeline.delete_adapters([qualname])
+            self.lora_adapters.remove(qualname)
+        else:
+            raise NotImplementedError(
+                f'Addon support for type {addon_type} not implemented'
+            )
+
+    def activate_addon(self, addon_account_id: str, addon_model_id: str):
+        qualname = f"accounts/{addon_account_id}/models/{addon_model_id}"
+
+        if qualname in self.lora_adapters:
+            if self.active_lora_adapter is not None:
+                raise ValueError(
+                    f"LoRA adapter {self.active_lora_adapter} already active. Multi-LoRA not yet supported"
+                )
+
+            self.active_lora_adapter = qualname
+            return
+
+        raise ValueError(f"Unknown addon {qualname}")
+
+    def deactivate_addon(self, addon_account_id: str, addon_model_id: str):
+        qualname = f"accounts/{addon_account_id}/models/{addon_model_id}"
+
+        if self.active_lora_adapter == qualname:
+            self.active_lora_adapter = None
+        else:
+            raise AssertionError(f'Addon {qualname} not loaded!')
+
+
+if __name__ == "__flumina_main__":
+    f = FluminaModule()
+    flumina_main(f)
+
+if __name__ == "__main__":
+    f = FluminaModule()
+    f._test_return_sync_response = True
+    import asyncio
+    out = asyncio.run(f.text_to_image(
+        body=Text2ImageRequest(
+            prompt="test"
+        ),
+        accept="image/png"
+    ))
+    with open("out_image.png", "wb") as f:
+        f.write(out)
+

requirements.txt

@@ -1,3 +1,6 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
+torch>=2.4
+torchvision
 git+https://github.com/aredden/torch-cublas-hgemm.git@master
 einops
 PyTurboJPEG
@@ -13,4 +16,4 @@ accelerate
 quanto
 pydash
 pybase64
-uvicorn
+uvicorn