<!-- livebook:{"app_settings":{"access_type":"public","auto_shutdown_ms":3600000,"multi_session":true,"show_existing_sessions":false,"slug":"stability"},"autosave_interval_s":null,"file_entries":[{"file":{"file_system_id":"local","file_system_type":"local","path":"/data/giraffe-explosion.webp"},"name":"giraffe-explosion.webp","type":"file"}]} -->

# Stability AI

```elixir
Mix.install([
  {:tesla, "~> 1.4"},

  # optional, but recommended adapter
  {:hackney, "~> 1.17"},

  # optional, required by JSON middleware
  {:jason, ">= 1.0.0"},
  {:kino, "~> 0.12.0"}
])

Kino.nothing()
```

## Generate images with Stable Diffusion

```elixir
Kino.Markdown.new("""
Stability AI provides their Stability Diffusion family of image generation models. This application allows you to interface with Stability AI's API:s to generate images with your choice of models.
  
*To be able to generate images you need your Stability AI API-key. When you provide the key it will be used with every request to Stability AI but otherwise will be private to your session - meaning you will be asked again for your API-key if you were to close this session or if it were timed out.*

""")
```

```elixir
#Kino.FS.file_path("giraffe-explosion.webp")
#|> File.read!()
#|> Kino.Image.new("image/webp")
#|> Kino.render()
#
#Kino.Markdown.new(
#  "> Example of image generated with Stable Diffusion 3: `A running giraffe, behind it an explosion, all surrounded by the vastness of space and a galaxy.`  "
#)

Kino.nothing()
```

```elixir
defmodule StabilityAI do
  use Tesla

  @sd3_uri "/v2beta/stable-image/generate/sd3"
  @models %{
    sd1_6: "stable-diffusion-v1-6",
    sdxl: "stable-diffusion-xl-1024-v1-0",
    sd3: "sd3",
    sd3_turbo: "sd3-turbo"
  }

  @sd_uris %{
    sd1_6: "/v1/generation/stable-diffusion-v1-6/text-to-image",
    sdxl: "/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image",
    sd3: "/v2beta/stable-image/generate/sd3",
    sd3_turbo: "/v2beta/stable-image/generate/sd3"
  }

  @base_url "https://api.stability.ai"

  def image_path() do
    Kino.tmp_dir() <> "/image_generations/"
  end

  def get_image(model_prompt, model, api_key) do
    case model do
      :sd1_6 -> get_image_sd1(model_prompt, model, api_key)
      :sdxl -> get_image_sd1(model_prompt, model, api_key)
      :sd3 -> get_image_sd3(model_prompt, model, api_key)
      :sd3_turbo -> get_image_sd3(model_prompt, model, api_key)
    end
  end

  defp get_image_sd1(
         %{
           prompt: prompt,
           negative_prompt: _negative_prompt,
           aspect_ratio: aspect_ratio,
           style_preset: style_preset
         },
         model,
         api_key
       )
       when byte_size(prompt) > 5 and byte_size(aspect_ratio) >= 3 do
    client =
      Tesla.client([
        {Tesla.Middleware.BaseUrl, @base_url},
        Tesla.Middleware.JSON,
        {Tesla.Middleware.Headers,
         [
           {"Authorization", "Bearer #{api_key}"},
           {"Accept", "image/png"},
           {"Stability-Client-ID", "Livebook Stability"},
           {"Stability-Client-Version", "0.1"}
         ]}
      ])

    uri = Map.get(@sd_uris, model)

    {width, height} =
      case aspect_ratio do
        "21:9" -> {1536, 640}
        "16:9" -> {1344, 768}
        "3:2" -> {1216, 832}
        "1:1" -> {1024, 1024}
        "2:3" -> {832, 1216}
        "9:16" -> {768, 1344}
        "9:21" -> {640, 1536}
      end

    body = %{
      text_prompts: [%{text: prompt, weight: 0.5}],
      height: height,
      width: width,
      cfg_scale: 7,
      clip_guidance_preset: "NONE",
      samples: 1,
      seed: 0,
      steps: 30
    }

    body =
      case style_preset do
        nil -> body
        _ -> Map.put(body, :style_preset, style_preset)
      end

    post(
      client,
      uri,
      body
    )
  end

  defp get_image_sd3(
         %{
           prompt: prompt,
           negative_prompt: negative_prompt,
           aspect_ratio: aspect_ratio
         },
         model,
         api_key
       )
       when byte_size(prompt) > 5 and byte_size(aspect_ratio) >= 3 do
    client =
      Tesla.client([
        {Tesla.Middleware.BaseUrl, @base_url},
        Tesla.Middleware.FormUrlencoded
      ])

    body_template =
      Tesla.Multipart.new()
      |> Tesla.Multipart.add_field("prompt", prompt)
      |> Tesla.Multipart.add_field("output_format", "png")
      |> Tesla.Multipart.add_field("aspect_ratio", aspect_ratio)
      |> Tesla.Multipart.add_field("model", Map.get(@models, model))

    body =
      case byte_size(negative_prompt) do
        0 -> body_template
        _ -> body_template |> Tesla.Multipart.add_field("negative_prompt", negative_prompt)
      end

    post(
      client,
      @sd3_uri,
      body,
      headers: [
        {"Authorization", "Bearer #{api_key}"},
        {"Accept", "image/*"}
      ]
    )
  end

  def get_balance(api_key) do
    client =
      Tesla.client([
        {Tesla.Middleware.BaseUrl, @base_url},
        Tesla.Middleware.JSON,
        {Tesla.Middleware.Headers,
         [
           {"Authorization", "Bearer #{api_key}"},
           {"Accept", "application/json"},
           {"Stability-Client-ID", "Livebook Stability"},
           {"Stability-Client-Version", "0.1"}
         ]}
      ])

    get(
      client,
      "/v1/user/balance"
    )
  end

  def throttled_submission(form, data, status_frame: stability_status) do
    case RateLimitedForm.rate_limit(form) do
      :ok ->
        IO.puts("Generating with Stability AI...")
        Kino.Frame.render(stability_status, Kino.Markdown.new("Loading..."))

        StabilityAI.get_image(data, ModelVault.value(), KeyVault.value())

      :halt ->
        IO.puts("Please wait a few seconds before submitting again")
        nil
    end
  end

  def handle_response(
        response,
        data,
        status_frame: stability_status,
        preview_frame: preview_frame,
        output_frame: output_frame,
        balance_frame: balance_frame
      ) do
    case response do
      {:ok, %Tesla.Env{status: 200, body: body}} ->
        filename_template =
          :crypto.hash(:sha, "#{data.prompt} - #{:os.system_time(:second)}")
          |> Base.encode16(case: :lower)

        filename = filename_template <> ".png"

        unless File.exists?(image_path()) do
          File.mkdir!(image_path())
        end

        filepath = image_path() <> filename

        File.write!(filepath, body)
        IO.puts("Written file to '#{filepath}'")

        content = File.read!(filepath)
        image = Kino.Image.new(content, "image/png")

        ImagesVault.store(filepath)

        Kino.Frame.render(stability_status, Kino.Markdown.new("Done!"))

        Kino.Frame.render(preview_frame, image)

        images =
          for image_path <- ImagesVault.value() do
            File.read!(image_path) |> Kino.Image.new(:png)
          end

        images_grid = Kino.Layout.grid(images, columns: 3)

        Kino.Frame.render(output_frame, images_grid)

        case StabilityAI.get_balance(KeyVault.value()) do
          {:ok, %Tesla.Env{status: 200, body: body}} ->
            %{"credits" => credits} = body
            Kino.Frame.render(balance_frame, Kino.Markdown.new("**Balance** `#{credits}`"))

          _ ->
            :error
        end

        {:ok, filepath}

      _ ->
        Kino.Frame.render(stability_status, Kino.Markdown.new("Something went wrong!"))
        IO.inspect(response)
        {:error, response}
    end
  end
end

defmodule RateLimitedForm do
  use GenServer

  def start_link(form, opts) do
    GenServer.start_link(__MODULE__, {form, opts}, name: __MODULE__)
  end

  def init({form, _opts}) do
    {:ok, %{form: form, last_submitted: nil}}
  end

  def rate_limit(form) do
    GenServer.call(__MODULE__, {:rate_limit, form})
  end

  def handle_call({:rate_limit, _form}, _from, state) do
    current_time = :os.system_time(:second)

    case state.last_submitted && current_time - state.last_submitted < 5 do
      true ->
        {:reply, :halt, state}

      _ ->
        state = %{state | last_submitted: current_time}
        {:reply, :ok, state}
    end
  end
end

defmodule KeyVault do
  use Agent

  def start_link(initial_value) do
    Agent.start_link(fn -> initial_value end, name: __MODULE__)
  end

  def value do
    Agent.get(__MODULE__, & &1)
  end

  def store(value) do
    Agent.update(__MODULE__, fn _ -> value end)
  end
end

defmodule ImagesVault do
  use Agent

  def start_link(initial_value) do
    Agent.start_link(fn -> initial_value end, name: __MODULE__)
  end

  def value do
    Agent.get(__MODULE__, & &1)
  end

  def latest do
    Agent.get(__MODULE__, fn images ->
      List.first(images)
    end)
  end

  def store(image) do
    Agent.update(__MODULE__, fn images -> [image | images] end)
  end
end

defmodule ModelVault do
  use Agent

  def start_link(initial_value) do
    Agent.start_link(fn -> initial_value end, name: __MODULE__)
  end

  def value do
    Agent.get(__MODULE__, & &1)
  end

  def store(value) do
    Agent.update(__MODULE__, fn _ -> value end)
  end
end

KeyVault.start_link("")
ImagesVault.start_link([])
ModelVault.start_link(:sd1_6)

balance_frame = Kino.Frame.new()

api_key_form =
  Kino.Control.form(
    [
      api_key: Kino.Input.password("Stability API Key")
    ],
    submit: "OK"
  )

api_key_status = Kino.Frame.new(placeholder: false)

Kino.listen(api_key_form, fn event ->
  case event do
    %{type: :submit} ->
      %{data: %{api_key: api_key}} = event

      KeyVault.store(api_key)
      Kino.Frame.append(api_key_status, Kino.Markdown.new("*Updated key*"))

      case StabilityAI.get_balance(KeyVault.value()) do
        {:ok, %Tesla.Env{status: 200, body: body}} ->
          %{"credits" => credits} = body
          Kino.Frame.render(balance_frame, Kino.Markdown.new("**Balance** `#{credits}`"))

        _ ->
          :error
      end

    _ ->
      IO.puts("Error: Not specified event for API Key form!")
  end
end)

basic_form_fields = [
  prompt: Kino.Input.textarea("Prompt"),
  negative_prompt: Kino.Input.text("Negative Prompt"),
  aspect_ratio:
    Kino.Input.select("Aspect Ratio", [
      {"21:9", "21:9"},
      {"16:9", "16:9"},
      {"3:2", "3:2"},
      {"1:1", "1:1"},
      {"2:3", "2:3"},
      {"9:16", "9:16"},
      {"9:21", "9:21"}
    ])
]

v1_form =
  Kino.Control.form(
    basic_form_fields ++
      [
        style_preset:
          Kino.Input.select("Predefined Style", [
            {"3d-model", "3D Model"},
            {"analog-film", "Analog Film"},
            {"anime", "Anime"},
            {"cinematic", "Cinematic"},
            {"comic-book", "Comic Book"},
            {"digital-art", "Digital Art"},
            {"enhance", "Enhance"},
            {"fantasy-art", "Fantasy Art"},
            {"isometric", "Isometric"},
            {"line-art", "Line Art"},
            {"low-poly", "Low Poly"},
            {"modeling-compound", "Modeling Compound"},
            {"neon-punk", "Neon Punk"},
            {"origami", "Origami"},
            {"photographic", "Photographic"},
            {"pixel-art", "Pixel Art"},
            {"tile-texture", "Tile Texture"},
            {nil, "None"}
          ])
      ],
    submit: "Submit"
  )

v3_form =
  Kino.Control.form(
    basic_form_fields,
    submit: "Submit"
  )

prompt_frame = Kino.Frame.new()
Kino.Frame.render(prompt_frame, v1_form)

case RateLimitedForm.start_link(v1_form, []) do
  {:error, {:already_started, _pid}} -> :ok
  {:ok, _pid} -> :ok
  x -> x
end

case RateLimitedForm.start_link(v3_form, []) do
  {:error, {:already_started, _pid}} -> :ok
  {:ok, _pid} -> :ok
  x -> x
end

Kino.nothing()
```

```elixir
Kino.Layout.grid(
  [
    api_key_form,
    api_key_status
  ],
  columns: 2
)
```

```elixir
# Select model to use 
# model_form =
#   Kino.Control.form(
#     [
#       model:
#         Kino.Input.select("Model", [
#           {:sd1_6, "SD 1.6"},
#           {:sdxl, "SDXL"},
#           {:sd3, "SD 3"},
#           {:sd3_turbo, "SD 3-Turbo"}
#         ])
#     ],
#     submit: "Select Model"
#   )

model_selection =
  Kino.Input.select("Model", [
    {:sd1_6, "SD 1.6"},
    {:sdxl, "SDXL"},
    {:sd3, "SD 3"},
    {:sd3_turbo, "SD 3-Turbo"}
  ])

Kino.listen(model_selection, fn event ->
  %{value: model} = event
  ModelVault.store(model)

  case model do
    :sd1_6 -> Kino.Frame.render(prompt_frame, v1_form)
    :sdxl -> Kino.Frame.render(prompt_frame, v1_form)
    :sd3 -> Kino.Frame.render(prompt_frame, v3_form)
    :sd3_turbo -> Kino.Frame.render(prompt_frame, v3_form)
  end
end)

model_selection
```

````elixir
stability_status = Kino.Frame.new()
prompt_status = Kino.Frame.new()

preview_frame = Kino.Frame.new(placeholder: false)
latest_image_file = ImagesVault.latest()

case latest_image_file do
  nil ->
    nil

  _ ->
    latest_image = File.read!(latest_image_file) |> Kino.Image.new(:png)
    Kino.Frame.render(preview_frame, latest_image)
end

status_layout =
  Kino.Layout.grid([Kino.Markdown.new("### *Status*"), prompt_status, stability_status],
    columns: 1
  )

form_frame = Kino.Frame.new()
form_layout = Kino.Layout.grid([form_frame, status_layout], columns: 2)

Kino.render(form_layout)

Kino.Frame.append(form_frame, Kino.Markdown.new("## Generate Image"))
Kino.Frame.append(form_frame, balance_frame)
Kino.Frame.append(form_frame, prompt_frame)

images =
  for image_path <- ImagesVault.value() do
    File.read!(image_path) |> Kino.Image.new(:png)
  end

images_grid = Kino.Layout.grid(images, columns: 3)
output_frame = Kino.Frame.new()

Kino.Markdown.new("""
  
___  
   
Preview:  
  
""")
|> Kino.render()

Kino.render(preview_frame)
Kino.Frame.append(output_frame, images_grid)

Kino.listen(v1_form, fn event ->
  case event do
    %{type: :submit} ->
      %{data: data} = event

      model = ModelVault.value()

      Kino.Frame.render(
        prompt_status,
        Kino.Markdown.new("""
        ```json
        {
          "prompt": #{data.prompt}  
          "negative_prompt": #{data.negative_prompt}  
          "aspect_ratio": #{data.aspect_ratio}  
          "model": #{model}
        }
        ```
        """)
      )

      response = StabilityAI.throttled_submission(v1_form, data, status_frame: stability_status)

      StabilityAI.handle_response(
        response,
        data,
        status_frame: stability_status,
        preview_frame: preview_frame,
        output_frame: output_frame,
        balance_frame: balance_frame
      )

    _ ->
      IO.puts("What?")
  end
end)

Kino.listen(v3_form, fn event ->
  case event do
    %{type: :submit} ->
      %{data: data} = event

      model = ModelVault.value()

      Kino.Frame.render(
        prompt_status,
        Kino.Markdown.new("""
        ```json
        {
          "prompt": #{data.prompt}  
          "negative_prompt": #{data.negative_prompt}  
          "aspect_ratio": #{data.aspect_ratio}  
          "model": #{model}
        }
        ```
        """)
      )

      response = StabilityAI.throttled_submission(v1_form, data, status_frame: stability_status)

      StabilityAI.handle_response(
        response,
        data,
        status_frame: stability_status,
        preview_frame: preview_frame,
        output_frame: output_frame,
        balance_frame: balance_frame
      )

    _ ->
      IO.puts("What?")
  end
end)

Kino.Markdown.new("""
  
___  
  

  
""")
|> Kino.render()

output_frame
````

```elixir
download_button = Kino.Control.button("ZIP images")
download_frame = Kino.Frame.new()
Kino.render(download_button)
Kino.render(download_frame)

Kino.listen(download_button, fn _event ->
  now =
    DateTime.utc_now()
    |> DateTime.to_string()
    |> String.replace(~r/[:\s-]/, "")
    |> String.replace(".", "_")

  image_files = ImagesVault.value()

  files_to_zip =
    image_files
    |> Enum.map(fn file_path ->
      String.replace(file_path, StabilityAI.image_path(), "")
      |> String.to_charlist()
    end)

  zip_filename = "stability_images_#{now}.zip"
  zip_filepath = StabilityAI.image_path() <> zip_filename

  case :zip.create(
         String.to_charlist(zip_filepath),
         files_to_zip,
         cwd: String.to_charlist(StabilityAI.image_path())
       ) do
    {:ok, _} ->
      download_prompt =
        Kino.Download.new(
          fn -> "#{zip_filepath}" |> File.read!() end,
          filename: zip_filename,
          label: zip_filename
        )

      IO.puts("ZIP file created: #{zip_filename}")

      Kino.Frame.render(download_frame, download_prompt)

    x ->
      IO.inspect(x)
      IO.puts("Error creating ZIP")
  end
end)

Kino.nothing()
```