unsmart / public-apps /stability.livemd
beercan's picture
Update public-apps/stability.livemd
3ad04e7 verified
raw
history blame
16.6 kB
<!-- 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()
```