Preset improvements

lib/config.py

@@ -1,6 +1,19 @@
 from dataclasses import dataclass
 from typing import Dict, List
 
+from .preset import preset
+
+
+def txt2img_models_from_presets(presets):
+    models = {}
+    for p in presets:
+        service = p.service
+        model_id = p.model_id
+        if service not in models:
+            models[service] = []
+        models[service].append(model_id)
+    return models
+
 
 @dataclass
 class Txt2TxtConfig:
@@ -32,6 +45,7 @@ class Config:
     txt2txt: Txt2TxtConfig
 
 
+# TODO: API keys should be with services (make a dataclass)
 config = Config(
     title="API Inference",
     icon="⚡",
@@ -50,34 +64,9 @@ config = Config(
             "Hugging Face": 2,
             "Together": 0,
         },
-        models={
-            # Model identifiers referenced in Text_to_Image.py
-            "Black Forest Labs": [
-                "flux-dev",
-                "flux-pro",
-                "flux-pro-1.1",
-            ],
-            "Fal": [
-                "fal-ai/aura-flow",
-                "fal-ai/flux/dev",
-                "fal-ai/flux/schnell",
-                "fal-ai/flux-pro",
-                "fal-ai/flux-pro/v1.1",
-                "fal-ai/fooocus",
-                "fal-ai/kolors",
-                "fal-ai/stable-diffusion-v3-medium",
-            ],
-            "Hugging Face": [
-                "black-forest-labs/flux.1-dev",
-                "black-forest-labs/flux.1-schnell",
-                "stabilityai/stable-diffusion-xl-base-1.0",
-            ],
-            "Together": [
-                "black-forest-labs/FLUX.1-schnell-Free",
-            ],
-        },
+        models=txt2img_models_from_presets(preset.txt2img.presets),
         hidden_parameters=[
-            # sent to API but not shown in generation parameters accordion
+            # Sent to API but not shown in generation parameters accordion
             "enable_safety_checker",
             "max_sequence_length",
             "num_images",
@@ -90,7 +79,7 @@ config = Config(
             "styles",
             "sync_mode",
         ],
-        negative_prompt="ugly, unattractive, disfigured, deformed, mutated, malformed, blurry, grainy, noisy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, autograph, trademark, sloppy, cluttered",
+        negative_prompt="ugly, unattractive, disfigured, deformed, mutated, malformed, blurry, grainy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, sloppy, cluttered",
         default_image_size="square_hd",
         image_sizes=[
             "landscape_16_9",
@@ -99,7 +88,7 @@ config = Config(
             "portrait_4_3",
             "portrait_16_9",
         ],
-        default_aspect_ratio="1024x1024",
+        default_aspect_ratio="1024x1024",  # fooocus aspect ratios
         aspect_ratios=[
             "704x1408",  # 1:2
             "704x1344",  # 11:21
@@ -143,6 +132,7 @@ config = Config(
                 "llama-3.1-sonar-large-128k-chat",
                 "llama-3.1-sonar-small-128k-online",
                 "llama-3.1-sonar-large-128k-online",
+                "llama-3.1-sonar-huge-128k-online",
             ],
         },
     ),

lib/preset.py

@@ -14,6 +14,8 @@ class Txt2TxtPreset:
 class Txt2ImgPreset:
     # FLUX1.1 has no scale or steps
     name: str
+    service: str
+    model_id: str
     guidance_scale: Optional[float] = None
     guidance_scale_min: Optional[float] = None
     guidance_scale_max: Optional[float] = None
@@ -32,25 +34,10 @@ class Txt2TxtPresets:
 
 @dataclass
 class Txt2ImgPresets:
-    # bfl
-    flux_1_1_pro_bfl: Txt2ImgPreset
-    flux_dev_bfl: Txt2ImgPreset
-    flux_pro_bfl: Txt2ImgPreset
-    # fal
-    aura_flow: Txt2ImgPreset
-    flux_1_1_pro_fal: Txt2ImgPreset
-    flux_dev_fal: Txt2ImgPreset
-    flux_pro_fal: Txt2ImgPreset
-    flux_schnell_fal: Txt2ImgPreset
-    fooocus: Txt2ImgPreset
-    kolors: Txt2ImgPreset
-    stable_diffusion_3: Txt2ImgPreset
-    # hf
-    flux_dev_hf: Txt2ImgPreset
-    flux_schnell_hf: Txt2ImgPreset
-    stable_diffusion_xl: Txt2ImgPreset
-    # together
-    flux_schnell_free_together: Txt2ImgPreset
+    presets: List[Txt2ImgPreset] = field(default_factory=list)
+
+    def __iter__(self):
+        return iter(self.presets)
 
 
 @dataclass
@@ -61,7 +48,6 @@ class Preset:
 
 preset = Preset(
     txt2txt=Txt2TxtPresets(
-        # Every service has model and system messages
         hugging_face=Txt2TxtPreset(
             frequency_penalty=0.0,
             frequency_penalty_min=-2.0,
@@ -76,177 +62,209 @@ preset = Preset(
         ),
     ),
     txt2img=Txt2ImgPresets(
-        aura_flow=Txt2ImgPreset(
-            "AuraFlow",
-            guidance_scale=3.5,
-            guidance_scale_min=1.0,
-            guidance_scale_max=10.0,
-            num_inference_steps=28,
-            num_inference_steps_min=10,
-            num_inference_steps_max=50,
-            parameters=["seed", "num_inference_steps", "guidance_scale", "expand_prompt"],
-            kwargs={"num_images": 1, "sync_mode": False},
-        ),
-        flux_1_1_pro_bfl=Txt2ImgPreset(
-            "FLUX1.1 Pro",
-            parameters=["seed", "width", "height", "prompt_upsampling"],
-            kwargs={"safety_tolerance": 6},
-        ),
-        flux_pro_bfl=Txt2ImgPreset(
-            "FLUX.1 Pro",
-            guidance_scale=2.5,
-            guidance_scale_min=1.5,
-            guidance_scale_max=5.0,
-            num_inference_steps=40,
-            num_inference_steps_min=10,
-            num_inference_steps_max=50,
-            parameters=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
-            kwargs={"safety_tolerance": 6, "interval": 1},
-        ),
-        flux_dev_bfl=Txt2ImgPreset(
-            "FLUX.1 Dev",
-            num_inference_steps=28,
-            num_inference_steps_min=10,
-            num_inference_steps_max=50,
-            guidance_scale=3.0,
-            guidance_scale_min=1.5,
-            guidance_scale_max=5.0,
-            parameters=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
-            kwargs={"safety_tolerance": 6},
-        ),
-        flux_1_1_pro_fal=Txt2ImgPreset(
-            "FLUX1.1 Pro",
-            parameters=["seed", "image_size"],
-            kwargs={
-                "num_images": 1,
-                "sync_mode": False,
-                "safety_tolerance": 6,
-                "enable_safety_checker": False,
-            },
-        ),
-        flux_pro_fal=Txt2ImgPreset(
-            "FLUX.1 Pro",
-            guidance_scale=2.5,
-            guidance_scale_min=1.5,
-            guidance_scale_max=5.0,
-            num_inference_steps=40,
-            num_inference_steps_min=10,
-            num_inference_steps_max=50,
-            parameters=["seed", "image_size", "num_inference_steps", "guidance_scale"],
-            kwargs={"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
-        ),
-        flux_dev_fal=Txt2ImgPreset(
-            "FLUX.1 Dev",
-            num_inference_steps=28,
-            num_inference_steps_min=10,
-            num_inference_steps_max=50,
-            guidance_scale=3.0,
-            guidance_scale_min=1.5,
-            guidance_scale_max=5.0,
-            parameters=["seed", "image_size", "num_inference_steps", "guidance_scale"],
-            kwargs={"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
-        ),
-        flux_schnell_fal=Txt2ImgPreset(
-            "FLUX.1 Schnell",
-            num_inference_steps=4,
-            num_inference_steps_min=1,
-            num_inference_steps_max=12,
-            parameters=["seed", "image_size", "num_inference_steps"],
-            kwargs={"num_images": 1, "sync_mode": False, "enable_safety_checker": False},
-        ),
-        flux_dev_hf=Txt2ImgPreset(
-            "FLUX.1 Dev",
-            num_inference_steps=28,
-            num_inference_steps_min=10,
-            num_inference_steps_max=50,
-            guidance_scale=3.0,
-            guidance_scale_min=1.5,
-            guidance_scale_max=5.0,
-            parameters=["width", "height", "guidance_scale", "num_inference_steps"],
-            kwargs={"max_sequence_length": 512},
-        ),
-        flux_schnell_hf=Txt2ImgPreset(
-            "FLUX.1 Schnell",
-            num_inference_steps=4,
-            num_inference_steps_min=1,
-            num_inference_steps_max=12,
-            parameters=["width", "height", "num_inference_steps"],
-            kwargs={"guidance_scale": 0.0, "max_sequence_length": 256},
-        ),
-        flux_schnell_free_together=Txt2ImgPreset(
-            "FLUX.1 Schnell Free",
-            num_inference_steps=4,
-            num_inference_steps_min=1,
-            num_inference_steps_max=12,
-            parameters=["model", "seed", "width", "height", "steps"],
-            kwargs={"n": 1},
-        ),
-        fooocus=Txt2ImgPreset(
-            "Fooocus",
-            guidance_scale=4.0,
-            guidance_scale_min=1.0,
-            guidance_scale_max=10.0,
-            parameters=["seed", "negative_prompt", "aspect_ratio", "guidance_scale"],
-            kwargs={
-                "num_images": 1,
-                "sync_mode": True,
-                "enable_safety_checker": False,
-                "output_format": "png",
-                "sharpness": 2,
-                "styles": ["Fooocus Enhance", "Fooocus V2", "Fooocus Sharp"],
-                "performance": "Quality",
-            },
-        ),
-        kolors=Txt2ImgPreset(
-            "Kolors",
-            guidance_scale=5.0,
-            guidance_scale_min=1.0,
-            guidance_scale_max=10.0,
-            num_inference_steps=50,
-            num_inference_steps_min=10,
-            num_inference_steps_max=50,
-            parameters=["seed", "negative_prompt", "image_size", "guidance_scale", "num_inference_steps"],
-            kwargs={
-                "num_images": 1,
-                "sync_mode": True,
-                "enable_safety_checker": False,
-                "scheduler": "EulerDiscreteScheduler",
-            },
-        ),
-        stable_diffusion_3=Txt2ImgPreset(
-            "SD3",
-            guidance_scale=5.0,
-            guidance_scale_min=1.0,
-            guidance_scale_max=10.0,
-            num_inference_steps=28,
-            num_inference_steps_min=10,
-            num_inference_steps_max=50,
-            parameters=[
-                "seed",
-                "negative_prompt",
-                "image_size",
-                "guidance_scale",
-                "num_inference_steps",
-                "prompt_expansion",
-            ],
-            kwargs={"num_images": 1, "sync_mode": True, "enable_safety_checker": False},
-        ),
-        stable_diffusion_xl=Txt2ImgPreset(
-            "SDXL",
-            guidance_scale=7.0,
-            guidance_scale_min=1.0,
-            guidance_scale_max=10.0,
-            num_inference_steps=40,
-            num_inference_steps_min=10,
-            num_inference_steps_max=50,
-            parameters=[
-                "seed",
-                "negative_prompt",
-                "width",
-                "height",
-                "guidance_scale",
-                "num_inference_steps",
-            ],
-        ),
+        presets=[
+            Txt2ImgPreset(
+                "AuraFlow",
+                "Fal",
+                "fal-ai/aura-flow",
+                guidance_scale=3.5,
+                guidance_scale_min=1.0,
+                guidance_scale_max=10.0,
+                num_inference_steps=28,
+                num_inference_steps_min=10,
+                num_inference_steps_max=50,
+                parameters=["seed", "num_inference_steps", "guidance_scale", "expand_prompt"],
+                kwargs={"num_images": 1, "sync_mode": False},
+            ),
+            Txt2ImgPreset(
+                "FLUX1.1 Pro",
+                "Black Forest Labs",
+                "flux-pro-1.1",
+                parameters=["seed", "width", "height", "prompt_upsampling"],
+                kwargs={"safety_tolerance": 6},
+            ),
+            Txt2ImgPreset(
+                "FLUX.1 Pro",
+                "Black Forest Labs",
+                "flux-pro",
+                guidance_scale=2.5,
+                guidance_scale_min=1.5,
+                guidance_scale_max=5.0,
+                num_inference_steps=40,
+                num_inference_steps_min=10,
+                num_inference_steps_max=50,
+                parameters=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
+                kwargs={"safety_tolerance": 6, "interval": 1},
+            ),
+            Txt2ImgPreset(
+                "FLUX.1 Dev",
+                "Black Forest Labs",
+                "flux-dev",
+                num_inference_steps=28,
+                num_inference_steps_min=10,
+                num_inference_steps_max=50,
+                guidance_scale=3.0,
+                guidance_scale_min=1.5,
+                guidance_scale_max=5.0,
+                parameters=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
+                kwargs={"safety_tolerance": 6},
+            ),
+            Txt2ImgPreset(
+                "FLUX1.1 Pro",
+                "Fal",
+                "fal-ai/flux-pro/v1.1",
+                parameters=["seed", "image_size"],
+                kwargs={
+                    "num_images": 1,
+                    "sync_mode": False,
+                    "safety_tolerance": 6,
+                    "enable_safety_checker": False,
+                },
+            ),
+            Txt2ImgPreset(
+                "FLUX.1 Pro",
+                "Fal",
+                "fal-ai/flux-pro",
+                guidance_scale=2.5,
+                guidance_scale_min=1.5,
+                guidance_scale_max=5.0,
+                num_inference_steps=40,
+                num_inference_steps_min=10,
+                num_inference_steps_max=50,
+                parameters=["seed", "image_size", "num_inference_steps", "guidance_scale"],
+                kwargs={"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
+            ),
+            Txt2ImgPreset(
+                "FLUX.1 Dev",
+                "Fal",
+                "fal-ai/flux/dev",
+                num_inference_steps=28,
+                num_inference_steps_min=10,
+                num_inference_steps_max=50,
+                guidance_scale=3.0,
+                guidance_scale_min=1.5,
+                guidance_scale_max=5.0,
+                parameters=["seed", "image_size", "num_inference_steps", "guidance_scale"],
+                kwargs={"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
+            ),
+            Txt2ImgPreset(
+                "FLUX.1 Schnell",
+                "Fal",
+                "fal-ai/flux/schnell",
+                num_inference_steps=4,
+                num_inference_steps_min=1,
+                num_inference_steps_max=12,
+                parameters=["seed", "image_size", "num_inference_steps"],
+                kwargs={"num_images": 1, "sync_mode": False, "enable_safety_checker": False},
+            ),
+            Txt2ImgPreset(
+                "FLUX.1 Dev",
+                "Hugging Face",
+                "black-forest-labs/flux.1-dev",
+                num_inference_steps=28,
+                num_inference_steps_min=10,
+                num_inference_steps_max=50,
+                guidance_scale=3.0,
+                guidance_scale_min=1.5,
+                guidance_scale_max=5.0,
+                parameters=["width", "height", "guidance_scale", "num_inference_steps"],
+                kwargs={"max_sequence_length": 512},
+            ),
+            Txt2ImgPreset(
+                "FLUX.1 Schnell",
+                "Hugging Face",
+                "black-forest-labs/flux.1-schnell",
+                num_inference_steps=4,
+                num_inference_steps_min=1,
+                num_inference_steps_max=12,
+                parameters=["width", "height", "num_inference_steps"],
+                kwargs={"guidance_scale": 0.0, "max_sequence_length": 256},
+            ),
+            Txt2ImgPreset(
+                "FLUX.1 Schnell Free",
+                "Together",
+                "black-forest-labs/FLUX.1-schnell-Free",
+                num_inference_steps=4,
+                num_inference_steps_min=1,
+                num_inference_steps_max=12,
+                parameters=["model", "seed", "width", "height", "steps"],
+                kwargs={"n": 1},
+            ),
+            Txt2ImgPreset(
+                "Fooocus",
+                "Fal",
+                "fal-ai/fooocus",
+                guidance_scale=4.0,
+                guidance_scale_min=1.0,
+                guidance_scale_max=10.0,
+                parameters=["seed", "negative_prompt", "aspect_ratio", "guidance_scale"],
+                kwargs={
+                    "num_images": 1,
+                    "sync_mode": True,
+                    "enable_safety_checker": False,
+                    "output_format": "png",
+                    "sharpness": 2,
+                    "styles": ["Fooocus Enhance", "Fooocus V2", "Fooocus Sharp"],
+                    "performance": "Quality",
+                },
+            ),
+            Txt2ImgPreset(
+                "Kolors",
+                "Fal",
+                "fal-ai/kolors",
+                guidance_scale=5.0,
+                guidance_scale_min=1.0,
+                guidance_scale_max=10.0,
+                num_inference_steps=50,
+                num_inference_steps_min=10,
+                num_inference_steps_max=50,
+                parameters=["seed", "negative_prompt", "image_size", "guidance_scale", "num_inference_steps"],
+                kwargs={
+                    "num_images": 1,
+                    "sync_mode": True,
+                    "enable_safety_checker": False,
+                    "scheduler": "EulerDiscreteScheduler",
+                },
+            ),
+            Txt2ImgPreset(
+                "SD3",
+                "Fal",
+                "fal-ai/stable-diffusion-v3-medium",
+                guidance_scale=5.0,
+                guidance_scale_min=1.0,
+                guidance_scale_max=10.0,
+                num_inference_steps=28,
+                num_inference_steps_min=10,
+                num_inference_steps_max=50,
+                parameters=[
+                    "seed",
+                    "negative_prompt",
+                    "image_size",
+                    "guidance_scale",
+                    "num_inference_steps",
+                    "prompt_expansion",
+                ],
+                kwargs={"num_images": 1, "sync_mode": True, "enable_safety_checker": False},
+            ),
+            Txt2ImgPreset(
+                "SDXL",
+                "Hugging Face",
+                "stabilityai/stable-diffusion-xl-base-1.0",
+                guidance_scale=7.0,
+                guidance_scale_min=1.0,
+                guidance_scale_max=10.0,
+                num_inference_steps=40,
+                num_inference_steps_min=10,
+                num_inference_steps_max=50,
+                parameters=[
+                    "seed",
+                    "negative_prompt",
+                    "width",
+                    "height",
+                    "guidance_scale",
+                    "num_inference_steps",
+                ],
+            ),
+        ]
     ),
 )

pages/1_💬_Text_Generation.py

@@ -43,7 +43,7 @@ st.logo("logo.png")
 st.sidebar.header("Settings")
 service = st.sidebar.selectbox(
     "Service",
-    options=["Hugging Face", "Perplexity"],
+    options=SERVICE_SESSION.keys(),
     index=0,
     disabled=st.session_state.running,
 )

pages/2_🎨_Text_to_Image.py

@@ -20,29 +20,9 @@ SESSION_TOKEN = {
     "api_key_together": os.environ.get("TOGETHER_API_KEY") or None,
 }
 
-# TODO: group by service so we can have models with the same name
-# Model IDs in lib/config.py
-PRESET_MODEL = {
-    # bfl
-    "flux-pro-1.1": preset.txt2img.flux_1_1_pro_bfl,
-    "flux-pro": preset.txt2img.flux_pro_bfl,
-    "flux-dev": preset.txt2img.flux_dev_bfl,
-    # fal
-    "fal-ai/aura-flow": preset.txt2img.aura_flow,
-    "fal-ai/flux/dev": preset.txt2img.flux_dev_fal,
-    "fal-ai/flux/schnell": preset.txt2img.flux_schnell_fal,
-    "fal-ai/flux-pro": preset.txt2img.flux_pro_fal,
-    "fal-ai/flux-pro/v1.1": preset.txt2img.flux_1_1_pro_fal,
-    "fal-ai/fooocus": preset.txt2img.fooocus,
-    "fal-ai/kolors": preset.txt2img.kolors,
-    "fal-ai/stable-diffusion-v3-medium": preset.txt2img.stable_diffusion_3,
-    # hf
-    "black-forest-labs/flux.1-dev": preset.txt2img.flux_dev_hf,
-    "black-forest-labs/flux.1-schnell": preset.txt2img.flux_schnell_hf,
-    "stabilityai/stable-diffusion-xl-base-1.0": preset.txt2img.stable_diffusion_xl,
-    # together
-    "black-forest-labs/FLUX.1-schnell-Free": preset.txt2img.flux_schnell_free_together,
-}
+PRESET_MODEL = {}
+for p in preset.txt2img.presets:
+    PRESET_MODEL[p.model_id] = p
 
 st.set_page_config(
     page_title=f"{config.title} | Text to Image",
@@ -108,8 +88,8 @@ st.html("""
 
 # Build parameters from preset by rendering the appropriate input widgets
 parameters = {}
-preset = PRESET_MODEL[model]
-for param in preset.parameters:
+model_preset = PRESET_MODEL[model]
+for param in model_preset.parameters:
     if param == "model":
         parameters[param] = model
     if param == "seed":
@@ -161,18 +141,18 @@ for param in preset.parameters:
     if param in ["guidance_scale", "guidance"]:
         parameters[param] = st.sidebar.slider(
             "Guidance Scale",
-            preset.guidance_scale_min,
-            preset.guidance_scale_max,
-            preset.guidance_scale,
+            model_preset.guidance_scale_min,
+            model_preset.guidance_scale_max,
+            model_preset.guidance_scale,
             0.1,
             disabled=st.session_state.running,
         )
     if param in ["num_inference_steps", "steps"]:
         parameters[param] = st.sidebar.slider(
             "Inference Steps",
-            preset.num_inference_steps_min,
-            preset.num_inference_steps_max,
-            preset.num_inference_steps,
+            model_preset.num_inference_steps_min,
+            model_preset.num_inference_steps_max,
+            model_preset.num_inference_steps,
             1,
             disabled=st.session_state.running,
         )
@@ -280,8 +260,8 @@ if prompt := st.chat_input(
 
     with st.chat_message("assistant"):
         with st.spinner("Running..."):
-            if preset.kwargs:
-                parameters.update(preset.kwargs)
+            if model_preset.kwargs:
+                parameters.update(model_preset.kwargs)
             session_key = f"api_key_{service.lower().replace(' ', '_')}"
             api_key = st.session_state[session_key] or SESSION_TOKEN[session_key]
             image = txt2img_generate(api_key, service, model, prompt, parameters)