Support image-to-image workflows

lib/__init__.py

@@ -1,7 +1,10 @@
 from .api import txt2img_generate, txt2txt_generate
 from .config import config
+from .util import base64_decode_image_data_url, base64_encode_image_file
 
 __all__ = [
+    "base64_decode_image_data_url",
+    "base64_encode_image_file",
     "config",
     "txt2img_generate",
     "txt2txt_generate",

lib/api.py

@@ -1,6 +1,5 @@
-import base64
-import io
 import time
+from io import BytesIO
 
 import httpx
 import streamlit as st
@@ -11,6 +10,7 @@ from openai import OpenAI
 from PIL import Image
 
 from .config import config
+from .util import base64_decode_image_data_url
 
 
 def txt2txt_generate(api_key, service, parameters, **kwargs):
@@ -72,6 +72,7 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
     try:
         timeout = config.timeout
         response = httpx.post(base_url, headers=headers, json=json, timeout=timeout)
+
         if response.status_code // 100 == 2:  # 2xx
             # BFL is async so we need to poll for result
             # https://api.bfl.ml/docs
@@ -91,7 +92,7 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
                             headers=headers,
                             timeout=timeout,
                         )
-                        return Image.open(io.BytesIO(image.content))
+                        return Image.open(BytesIO(image.content))
 
                     retries += 1
                     time.sleep(1)
@@ -100,21 +101,20 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
 
             if service == "fal":
                 # Sync mode means wait for image base64 string instead of CDN link
+                url = response.json()["images"][0]["url"]
                 if parameters.get("sync_mode", True):
-                    bytes = base64.b64decode(response.json()["images"][0]["url"].split(",")[-1])
-                    return Image.open(io.BytesIO(bytes))
+                    return base64_decode_image_data_url(url)
                 else:
-                    url = response.json()["images"][0]["url"]
                     image = httpx.get(url, headers=headers, timeout=timeout)
-                    return Image.open(io.BytesIO(image.content))
+                    return Image.open(BytesIO(image.content))
 
             if service == "hf":
-                return Image.open(io.BytesIO(response.content))
+                return Image.open(BytesIO(response.content))
 
             if service == "together":
                 url = response.json()["data"][0]["url"]
                 image = httpx.get(url, headers=headers, timeout=timeout)
-                return Image.open(io.BytesIO(image.content))
+                return Image.open(BytesIO(image.content))
 
         else:
             return f"Error: {response.status_code} {response.text}"

lib/config.py

@@ -6,7 +6,9 @@ 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"
 
-IMAGE_IMAGE_SIZES = [
+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",
@@ -40,6 +42,8 @@ IMAGE_ASPECT_RATIOS = [
 
 IMAGE_RANGE = (256, 1408)
 
+STRENGTH_RANGE = (0.0, 1.0)
+
 
 @dataclass
 class ModelConfig:
@@ -68,6 +72,8 @@ class ImageModelConfig(ModelConfig):
     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
@@ -149,6 +155,7 @@ config = AppConfig(
     hidden_parameters=[
         # Sent to API but not shown in generation parameters accordion
         "enable_safety_checker",
+        "image_url",
         "max_sequence_length",
         "n",
         "num_images",
@@ -223,17 +230,69 @@ config = AppConfig(
                 "fal-ai/aura-flow": ImageModelConfig(
                     "AuraFlow",
                     guidance_scale=3.5,
-                    guidance_scale_range=(1.0, 10.0),
-                    num_inference_steps=28,
-                    num_inference_steps_range=(10, 50),
+                    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_IMAGE_SIZES,
+                    image_sizes=IMAGE_SIZES,
                     kwargs={
                         "num_images": 1,
                         "sync_mode": False,
@@ -244,7 +303,7 @@ config = AppConfig(
                 "fal-ai/flux-pro": ImageModelConfig(
                     "FLUX.1 Pro",
                     image_size="square_hd",
-                    image_sizes=IMAGE_IMAGE_SIZES,
+                    image_sizes=IMAGE_SIZES,
                     guidance_scale=2.5,
                     guidance_scale_range=(1.5, 5.0),
                     num_inference_steps=40,
@@ -255,18 +314,38 @@ config = AppConfig(
                 "fal-ai/flux/dev": ImageModelConfig(
                     "FLUX.1 Dev",
                     image_size="square_hd",
-                    image_sizes=IMAGE_IMAGE_SIZES,
+                    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, "safety_tolerance": 6},
+                    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_IMAGE_SIZES,
+                    image_sizes=IMAGE_SIZES,
                     num_inference_steps=4,
                     num_inference_steps_range=(1, 12),
                     parameters=["seed", "image_size", "num_inference_steps"],
@@ -274,6 +353,7 @@ config = AppConfig(
                 ),
                 "fal-ai/fooocus": ImageModelConfig(
                     "Fooocus",
+                    negative_prompt=FOOOCUS_NEGATIVE_PROMPT,
                     aspect_ratio="1024x1024",
                     aspect_ratios=IMAGE_ASPECT_RATIOS,
                     guidance_scale=4.0,
@@ -292,8 +372,9 @@ config = AppConfig(
                 ),
                 "fal-ai/kolors": ImageModelConfig(
                     "Kolors",
+                    negative_prompt=IMAGE_NEGATIVE_PROMPT,
                     image_size="square_hd",
-                    image_sizes=IMAGE_IMAGE_SIZES,
+                    image_sizes=IMAGE_SIZES,
                     guidance_scale=5.0,
                     guidance_scale_range=(1.0, 10.0),
                     num_inference_steps=50,
@@ -313,9 +394,9 @@ config = AppConfig(
                     },
                 ),
                 "fal-ai/stable-diffusion-v3-medium": ImageModelConfig(
-                    "SD3",
+                    "SD3 Medium",
                     image_size="square_hd",
-                    image_sizes=IMAGE_IMAGE_SIZES,
+                    image_sizes=IMAGE_SIZES,
                     guidance_scale=5.0,
                     guidance_scale_range=(1.0, 10.0),
                     num_inference_steps=28,

lib/util.py

@@ -0,0 +1,20 @@
+import mimetypes
+from base64 import b64decode, b64encode
+from io import BytesIO
+
+from PIL import Image
+
+
+def base64_decode_image_data_url(data_url: str) -> Image:
+    _, data = data_url.split("base64,", maxsplit=1)
+    byte_data = b64decode(data)
+    return Image.open(BytesIO(byte_data))
+
+
+def base64_encode_image_file(image_file: BytesIO) -> str:
+    file_type = image_file.type
+    if not file_type:
+        file_type = mimetypes.guess_type(image_file.name)[0]
+    file_data = image_file.read()
+    b64 = b64encode(file_data).decode("utf-8")
+    return f"data:{file_type};base64,{b64}"

pages/2_🎨_Text_to_Image.py

@@ -2,7 +2,7 @@ from datetime import datetime
 
 import streamlit as st
 
-from lib import config, txt2img_generate
+from lib import base64_encode_image_file, config, txt2img_generate
 
 st.set_page_config(
     page_title=f"Text to Image - {config.title}",
@@ -81,6 +81,7 @@ parameters = {}
 for param in model_config.parameters:
     if param == "model":
         parameters[param] = model
+
     if param == "seed":
         parameters[param] = st.sidebar.number_input(
             "Seed",
@@ -89,12 +90,14 @@ for param in model_config.parameters:
             value=-1,
             disabled=st.session_state.running,
         )
+
     if param == "negative_prompt":
         parameters[param] = st.sidebar.text_area(
             "Negative Prompt",
             value=model_config.negative_prompt,
             disabled=st.session_state.running,
         )
+
     if param == "width":
         parameters[param] = st.sidebar.slider(
             "Width",
@@ -104,6 +107,7 @@ for param in model_config.parameters:
             max_value=model_config.width_range[1],
             disabled=st.session_state.running,
         )
+
     if param == "height":
         parameters[param] = st.sidebar.slider(
             "Height",
@@ -113,6 +117,7 @@ for param in model_config.parameters:
             max_value=model_config.height_range[1],
             disabled=st.session_state.running,
         )
+
     if param == "image_size":
         parameters[param] = st.sidebar.select_slider(
             "Image Size",
@@ -120,6 +125,7 @@ for param in model_config.parameters:
             value=model_config.image_size,
             disabled=st.session_state.running,
         )
+
     if param == "aspect_ratio":
         parameters[param] = st.sidebar.select_slider(
             "Aspect Ratio",
@@ -127,6 +133,7 @@ for param in model_config.parameters:
             value=model_config.aspect_ratio,
             disabled=st.session_state.running,
         )
+
     if param in ["guidance_scale", "guidance"]:
         parameters[param] = st.sidebar.slider(
             "Guidance Scale",
@@ -136,6 +143,7 @@ for param in model_config.parameters:
             0.1,
             disabled=st.session_state.running,
         )
+
     if param in ["num_inference_steps", "steps"]:
         parameters[param] = st.sidebar.slider(
             "Inference Steps",
@@ -145,12 +153,24 @@ for param in model_config.parameters:
             1,
             disabled=st.session_state.running,
         )
+
+    if param == "strength":
+        parameters[param] = st.sidebar.slider(
+            "Strength",
+            model_config.strength_range[0],
+            model_config.strength_range[1],
+            model_config.strength,
+            0.05,
+            disabled=st.session_state.running,
+        )
+
     if param in ["expand_prompt", "prompt_expansion"]:
         parameters[param] = st.sidebar.checkbox(
             "Prompt Expansion",
             value=False,
             disabled=st.session_state.running,
         )
+
     if param == "prompt_upsampling":
         parameters[param] = st.sidebar.checkbox(
             "Prompt Upsampling",
@@ -158,6 +178,16 @@ for param in model_config.parameters:
             disabled=st.session_state.running,
         )
 
+    if param == "image_url":
+        image_file = st.sidebar.file_uploader(
+            "Image",
+            type=["bmp", "gif", "jpg", "jpeg", "png", "webp"],
+            accept_multiple_files=False,
+            disabled=st.session_state.running,
+        )
+        if image_file:
+            parameters[param] = base64_encode_image_file(image_file)
+
 # Wrap the prompt in an accordion to display additional parameters
 for message in st.session_state.txt2img_messages:
     role = message["role"]