From e9edd919bee125c966287a7e69c69febdd40d1d5 Mon Sep 17 00:00:00 2001
From: Anthony Wu <462072+anthonywu@users.noreply.github.com>
Date: Fri, 25 Oct 2024 11:36:35 -0700
Subject: [PATCH] support HF models - freepik as first example

---
 src/mflux/config/model_config.py | 4 ++++
 src/mflux/ui/defaults.py         | 3 ++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/mflux/config/model_config.py b/src/mflux/config/model_config.py
index e67a581..56e793f 100644
--- a/src/mflux/config/model_config.py
+++ b/src/mflux/config/model_config.py
@@ -2,8 +2,12 @@
 
 
 class ModelConfig(Enum):
+    # ==== black-forest-labs official models =====
     FLUX1_DEV = ("black-forest-labs/FLUX.1-dev", "dev", 1000, 512)
     FLUX1_SCHNELL = ("black-forest-labs/FLUX.1-schnell", "schnell", 1000, 256)
+    # ==== third party compatible models - best effort compatibility =====
+    # https://huggingface.co/Freepik/flux.1-lite-8B-alpha distilled from dev
+    FLUX1_FREEPIK_LITE_8B_ALPHA = ("Freepik/flux.1-lite-8B-alpha", "freepik-lite-8b-alpha", 1000, 512)
 
     def __init__(
         self,
diff --git a/src/mflux/ui/defaults.py b/src/mflux/ui/defaults.py
index 7ed4bd7..a347231 100644
--- a/src/mflux/ui/defaults.py
+++ b/src/mflux/ui/defaults.py
@@ -2,9 +2,10 @@
 GUIDANCE_SCALE = 3.5
 HEIGHT, WIDTH = 1024, 1024
 INIT_IMAGE_STRENGTH = 0.4  # for image-to-image init_image
-MODEL_CHOICES = ["dev", "schnell"]
+MODEL_CHOICES = ["dev", "schnell", "freepik-lite-8b-alpha"]
 MODEL_INFERENCE_STEPS = {
     "dev": 14,
     "schnell": 4,
+    "freepik-lite-8b-alpha": 22,
 }
 QUANTIZE_CHOICES = [4, 8]