import os from dataclasses import dataclass, field from typing import Dict, List, Optional, Union TEXT_SYSTEM_PROMPT = "You are a helpful assistant. Be precise and concise." IMAGE_NEGATIVE_PROMPT = "ugly, unattractive, disfigured, deformed, mutated, malformed, blurry, grainy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, sloppy, cluttered" FOOOCUS_NEGATIVE_PROMPT = "(worst quality, low quality, normal quality, lowres, low details, oversaturated, undersaturated, overexposed, underexposed, grayscale, bw, bad photo, bad photography, bad art:1.4), (watermark, signature, text font, username, error, logo, words, letters, digits, autograph, trademark, name:1.2), (blur, blurry, grainy), morbid, ugly, asymmetrical, mutated malformed, mutilated, poorly lit, bad shadow, draft, cropped, out of frame, cut off, censored, jpeg artifacts, out of focus, glitch, duplicate, (airbrushed, cartoon, anime, semi-realistic, cgi, render, blender, digital art, manga, amateur:1.3), (3D ,3D Game, 3D Game Scene, 3D Character:1.1), (bad hands, bad anatomy, bad body, bad face, bad teeth, bad arms, bad legs, deformities:1.3)" IMAGE_SIZES = [ "landscape_16_9", "landscape_4_3", "square_hd", "portrait_4_3", "portrait_16_9", ] IMAGE_ASPECT_RATIOS = [ "704x1408", # 1:2 "704x1344", # 11:21 "768x1344", # 4:7 "768x1280", # 3:5 "832x1216", # 13:19 "832x1152", # 13:18 "896x1152", # 7:9 "896x1088", # 14:17 "960x1088", # 15:17 "960x1024", # 15:16 "1024x1024", "1024x960", # 16:15 "1088x960", # 17:15 "1088x896", # 17:14 "1152x896", # 9:7 "1152x832", # 18:13 "1216x832", # 19:13 "1280x768", # 5:3 "1344x768", # 7:4 "1344x704", # 21:11 "1408x704", # 2:1 ] IMAGE_RANGE = (256, 1408) STRENGTH_RANGE = (0.0, 1.0) @dataclass class ModelConfig: name: str parameters: List[str] kwargs: Optional[Dict[str, Union[str, int, float, bool]]] = field(default_factory=dict) @dataclass class TextModelConfig(ModelConfig): system_prompt: Optional[str] = None frequency_penalty: Optional[float] = None frequency_penalty_range: Optional[tuple[float, float]] = None presence_penalty: Optional[float] = None presence_penalty_range: Optional[tuple[float, float]] = None max_tokens: Optional[int] = None max_tokens_range: Optional[tuple[int, int]] = None temperature: Optional[float] = None temperature_range: Optional[tuple[float, float]] = None @dataclass class ImageModelConfig(ModelConfig): negative_prompt: Optional[str] = None width: Optional[int] = None width_range: Optional[tuple[int, int]] = None height: Optional[int] = None height_range: Optional[tuple[int, int]] = None strength: Optional[float] = None strength_range: Optional[tuple[float, float]] = None image_size: Optional[str] = None image_sizes: Optional[List[str]] = field(default_factory=list) aspect_ratio: Optional[str] = None aspect_ratios: Optional[List[str]] = field(default_factory=list) guidance_scale: Optional[float] = None guidance_scale_range: Optional[tuple[float, float]] = None num_inference_steps: Optional[int] = None num_inference_steps_range: Optional[tuple[int, int]] = None @dataclass class ServiceConfig: name: str url: str api_key: Optional[str] text: Optional[Dict[str, TextModelConfig]] = field(default_factory=dict) image: Optional[Dict[str, ImageModelConfig]] = field(default_factory=dict) @dataclass class AppConfig: title: str layout: str logo: str icon: str timeout: int hidden_parameters: List[str] services: Dict[str, ServiceConfig] _anthropic_text_kwargs = { "system_prompt": TEXT_SYSTEM_PROMPT, "max_tokens": 512, "max_tokens_range": (512, 4096), "temperature": 0.5, "temperature_range": (0.0, 1.0), "parameters": ["max_tokens", "temperature"], } _hf_text_kwargs = { "system_prompt": TEXT_SYSTEM_PROMPT, "frequency_penalty": 0.0, "frequency_penalty_range": (-2.0, 2.0), "max_tokens": 512, "max_tokens_range": (512, 4096), "temperature": 1.0, "temperature_range": (0.0, 2.0), "parameters": ["max_tokens", "temperature", "frequency_penalty", "seed"], } _openai_text_kwargs = { "system_prompt": TEXT_SYSTEM_PROMPT, "frequency_penalty": 0.0, "frequency_penalty_range": (-2.0, 2.0), "presence_penalty": 0.0, "presence_penalty_range": (-2.0, 2.0), "max_tokens": 512, "max_tokens_range": (512, 4096), "temperature": 1.0, "temperature_range": (0.0, 2.0), "parameters": ["max_tokens", "temperature", "frequency_penalty", "presence_penalty", "seed"], } _pplx_text_kwargs = { "system_prompt": TEXT_SYSTEM_PROMPT, "frequency_penalty": 1.0, "frequency_penalty_range": (1.0, 2.0), "max_tokens": 512, "max_tokens_range": (512, 4096), "temperature": 1.0, "temperature_range": (0.0, 2.0), "parameters": ["max_tokens", "temperature", "frequency_penalty"], } config = AppConfig( title="Playground", layout="wide", logo="logo.svg", icon="⚡", timeout=60, hidden_parameters=[ # Sent to API but not shown in generation parameters accordion "enable_safety_checker", "image_url", "max_sequence_length", "n", "num_images", "output_format", "performance", "safety_tolerance", "scheduler", "sharpness", "style", "styles", "sync_mode", ], services={ "anthropic": ServiceConfig( name="Anthropic", url="https://api.anthropic.com/v1", api_key=os.environ.get("ANTHROPIC_API_KEY"), text={ "claude-3-haiku-20240307": TextModelConfig("Claude 3 Haiku", **_anthropic_text_kwargs), "claude-3-opus-20240229": TextModelConfig("Claude 3 Opus", **_anthropic_text_kwargs), "claude-3-sonnet-20240229": TextModelConfig("Claude 3 Sonnet", **_anthropic_text_kwargs), "claude-3-5-sonnet-20240620": TextModelConfig("Claude 3.5 Sonnet", **_anthropic_text_kwargs), }, ), "bfl": ServiceConfig( name="Black Forest Labs", url="https://api.bfl.ml/v1", api_key=os.environ.get("BFL_API_KEY"), image={ "flux-pro-1.1": ImageModelConfig( "FLUX1.1 Pro", width=1024, width_range=IMAGE_RANGE, height=1024, height_range=IMAGE_RANGE, parameters=["seed", "width", "height", "prompt_upsampling"], kwargs={"safety_tolerance": 6}, ), "flux-pro": ImageModelConfig( "FLUX.1 Pro", width=1024, width_range=IMAGE_RANGE, height=1024, height_range=IMAGE_RANGE, guidance_scale=2.5, guidance_scale_range=(1.5, 5.0), num_inference_steps=50, num_inference_steps_range=(10, 50), parameters=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"], kwargs={"safety_tolerance": 6, "interval": 1}, ), "flux-dev": ImageModelConfig( "FLUX.1 Dev", width=1024, width_range=IMAGE_RANGE, height=1024, height_range=IMAGE_RANGE, num_inference_steps=28, num_inference_steps_range=(10, 50), guidance_scale=3.0, guidance_scale_range=(1.5, 5.0), parameters=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"], kwargs={"safety_tolerance": 6}, ), }, ), "fal": ServiceConfig( name="Fal", url="https://fal.run", api_key=os.environ.get("FAL_KEY"), image={ "fal-ai/aura-flow": ImageModelConfig( "AuraFlow", guidance_scale=3.5, guidance_scale_range=(0.0, 20.0), num_inference_steps=50, num_inference_steps_range=(20, 50), parameters=["seed", "num_inference_steps", "guidance_scale", "expand_prompt"], kwargs={"num_images": 1, "sync_mode": False}, ), "fal-ai/fast-sdxl": ImageModelConfig( "Fast SDXL", negative_prompt=IMAGE_NEGATIVE_PROMPT, image_size="square_hd", image_sizes=IMAGE_SIZES, guidance_scale=7.5, guidance_scale_range=(0.0, 20.0), num_inference_steps=25, num_inference_steps_range=(1, 50), parameters=[ "seed", "negative_prompt", "image_size", "num_inference_steps", "guidance_scale", "expand_prompt", ], kwargs={ "num_images": 1, "sync_mode": False, "enable_safety_checker": False, "output_format": "png", }, ), "fal-ai/fast-sdxl/image-to-image": ImageModelConfig( "Fast SDXL (Image)", negative_prompt=IMAGE_NEGATIVE_PROMPT, image_size="square_hd", image_sizes=IMAGE_SIZES, strength=0.95, strength_range=STRENGTH_RANGE, guidance_scale=7.5, guidance_scale_range=(0.0, 20.0), num_inference_steps=25, num_inference_steps_range=(1, 50), parameters=[ "seed", "negative_prompt", "image_size", "num_inference_steps", "guidance_scale", "strength", "expand_prompt", "image_url", ], kwargs={ "num_images": 1, "sync_mode": False, "enable_safety_checker": False, "output_format": "png", }, ), "fal-ai/flux-pro/v1.1": ImageModelConfig( "FLUX1.1 Pro", parameters=["seed", "image_size"], image_size="square_hd", image_sizes=IMAGE_SIZES, kwargs={ "num_images": 1, "sync_mode": False, "safety_tolerance": 6, "enable_safety_checker": False, }, ), "fal-ai/flux-pro": ImageModelConfig( "FLUX.1 Pro", image_size="square_hd", image_sizes=IMAGE_SIZES, guidance_scale=2.5, guidance_scale_range=(1.5, 5.0), num_inference_steps=40, num_inference_steps_range=(10, 50), parameters=["seed", "image_size", "num_inference_steps", "guidance_scale"], kwargs={"num_images": 1, "sync_mode": False, "safety_tolerance": 6}, ), "fal-ai/flux/dev": ImageModelConfig( "FLUX.1 Dev", image_size="square_hd", image_sizes=IMAGE_SIZES, num_inference_steps=28, num_inference_steps_range=(10, 50), guidance_scale=3.0, guidance_scale_range=(1.5, 5.0), parameters=["seed", "image_size", "num_inference_steps", "guidance_scale"], kwargs={"num_images": 1, "sync_mode": False, "enable_safety_checker": False}, ), "fal-ai/flux/dev/image-to-image": ImageModelConfig( "FLUX.1 Dev (Image)", image_size="square_hd", image_sizes=IMAGE_SIZES, strength=0.95, strength_range=STRENGTH_RANGE, num_inference_steps=28, num_inference_steps_range=(10, 50), guidance_scale=3.0, guidance_scale_range=(1.5, 5.0), parameters=[ "seed", "image_size", "num_inference_steps", "guidance_scale", "strength", "image_url", ], kwargs={"num_images": 1, "sync_mode": False, "enable_safety_checker": False}, ), "fal-ai/flux/schnell": ImageModelConfig( "FLUX.1 Schnell", image_size="square_hd", image_sizes=IMAGE_SIZES, num_inference_steps=4, num_inference_steps_range=(1, 12), parameters=["seed", "image_size", "num_inference_steps"], kwargs={"num_images": 1, "sync_mode": False, "enable_safety_checker": False}, ), "fal-ai/fooocus": ImageModelConfig( "Fooocus", negative_prompt=FOOOCUS_NEGATIVE_PROMPT, aspect_ratio="1024x1024", aspect_ratios=IMAGE_ASPECT_RATIOS, guidance_scale=4.0, guidance_scale_range=(1.0, 15.0), parameters=["seed", "negative_prompt", "aspect_ratio", "guidance_scale"], # TODO: more of these can be params 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", }, ), "fal-ai/kolors": ImageModelConfig( "Kolors", negative_prompt=IMAGE_NEGATIVE_PROMPT, image_size="square_hd", image_sizes=IMAGE_SIZES, guidance_scale=5.0, guidance_scale_range=(1.0, 10.0), num_inference_steps=50, num_inference_steps_range=(10, 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", }, ), "fal-ai/stable-diffusion-v3-medium": ImageModelConfig( "SD3 Medium", image_size="square_hd", image_sizes=IMAGE_SIZES, guidance_scale=5.0, guidance_scale_range=(1.0, 10.0), num_inference_steps=28, num_inference_steps_range=(10, 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}, ), }, ), "hf": ServiceConfig( name="Hugging Face", url="https://api-inference.huggingface.co/models", api_key=os.environ.get("HF_TOKEN"), text={ "codellama/codellama-34b-instruct-hf": TextModelConfig("Code Llama 34B", **_hf_text_kwargs), "meta-llama/llama-2-13b-chat-hf": TextModelConfig("Meta Llama 2 13B", **_hf_text_kwargs), "mistralai/mistral-7b-instruct-v0.2": TextModelConfig("Mistral 0.2 7B", **_hf_text_kwargs), "nousresearch/nous-hermes-2-mixtral-8x7b-dpo": TextModelConfig( "Nous Hermes 2 Mixtral 8x7B", **_hf_text_kwargs, ), }, image={ "black-forest-labs/flux.1-dev": ImageModelConfig( "FLUX.1 Dev", width=1024, width_range=IMAGE_RANGE, height=1024, height_range=IMAGE_RANGE, guidance_scale=3.0, guidance_scale_range=(1.5, 5.0), num_inference_steps=28, num_inference_steps_range=(10, 50), parameters=["width", "height", "guidance_scale", "num_inference_steps"], ), "black-forest-labs/flux.1-schnell": ImageModelConfig( "FLUX.1 Schnell", width=1024, width_range=IMAGE_RANGE, height=1024, height_range=IMAGE_RANGE, num_inference_steps=4, num_inference_steps_range=(1, 12), parameters=["width", "height", "num_inference_steps"], kwargs={"guidance_scale": 0.0, "max_sequence_length": 256}, ), "stabilityai/stable-diffusion-xl-base-1.0": ImageModelConfig( "Stable Diffusion XL 1.0", negative_prompt=IMAGE_NEGATIVE_PROMPT, width=1024, width_range=IMAGE_RANGE, height=1024, height_range=IMAGE_RANGE, guidance_scale=7.0, guidance_scale_range=(1.0, 15.0), num_inference_steps=40, num_inference_steps_range=(10, 50), parameters=[ "seed", "negative_prompt", "width", "height", "guidance_scale", "num_inference_steps", ], ), }, ), "openai": ServiceConfig( name="OpenAI", url="https://api.openai.com/v1", api_key=os.environ.get("OPENAI_API_KEY"), text={ "chatgpt-4o-latest": TextModelConfig("ChatGPT-4o", **_openai_text_kwargs), "gpt-3.5-turbo": TextModelConfig("GPT-3.5 Turbo", **_openai_text_kwargs), "gpt-4-turbo": TextModelConfig("GPT-4 Turbo", **_openai_text_kwargs), "gpt-4o": TextModelConfig("GPT-4o", **_openai_text_kwargs), "gpt-4o-mini": TextModelConfig("GPT-4o mini", **_openai_text_kwargs), "o1-preview": TextModelConfig("o1-preview", **_openai_text_kwargs), "o1-mini": TextModelConfig("o1-mini", **_openai_text_kwargs), }, ), "pplx": ServiceConfig( name="Perplexity", url="https://api.perplexity.ai", api_key=os.environ.get("PPLX_API_KEY"), text={ "llama-3.1-sonar-small-128k-chat": TextModelConfig( "Sonar Small (Offline)", **_pplx_text_kwargs, ), "llama-3.1-sonar-large-128k-chat": TextModelConfig( "Sonar Large (Offline)", **_pplx_text_kwargs, ), "llama-3.1-sonar-small-128k-online": TextModelConfig( "Sonar Small (Online)", **_pplx_text_kwargs, ), "llama-3.1-sonar-large-128k-online": TextModelConfig( "Sonar Large (Online)", **_pplx_text_kwargs, ), "llama-3.1-sonar-huge-128k-online": TextModelConfig( "Sonar Huge (Online)", **_pplx_text_kwargs, ), }, ), # TODO: text models, more image models "together": ServiceConfig( name="Together", url="https://api.together.xyz/v1/images/generations", api_key=os.environ.get("TOGETHER_API_KEY"), image={ "black-forest-labs/FLUX.1-schnell-Free": ImageModelConfig( "FLUX.1 Schnell Free", width=1024, width_range=IMAGE_RANGE, height=1024, height_range=IMAGE_RANGE, num_inference_steps=4, num_inference_steps_range=(1, 12), parameters=["model", "seed", "width", "height", "steps"], kwargs={"n": 1}, ), }, ), }, )