progress_ipywidgets_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Stable Diffusion Demo\n",
    "## _with preview images_\n",
    "\n",
    "Stable Diffusion is a state of the art text-to-image model that generates images from text.\n",
    "\n",
    "For faster generation and forthcoming API access you can try [DreamStudio Beta](http://beta.dreamstudio.ai/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from datasets import load_dataset\n",
    "from PIL import Image  \n",
    "import re\n",
    "\n",
    "from preview_decoder import ApproximateDecoder, jpeg_bytes\n",
    "from generator_pipeline import StableDiffusionGeneratorPipeline, PipelineIntermediateState\n",
    "\n",
    "\n",
    "model_id = \"CompVis/stable-diffusion-v1-4\"\n",
    "device = \"cuda\"\n",
    "\n",
    "#If you are running this code locally, you need to either do a 'huggingface-cli login` or paste your User Access Token from here https://huggingface.co/settings/tokens into the use_auth_token field below. \n",
    "pipe = StableDiffusionGeneratorPipeline.from_pretrained(\n",
    "    model_id, use_auth_token=True, \n",
    "    revision=\"fp16\", torch_dtype=torch.float16,\n",
    ")\n",
    "pipe = pipe.to(device)\n",
    "pipe.enable_attention_slicing()\n",
    "torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "#When running locally, you won`t have access to this, so you can remove this part\n",
    "#word_list_dataset = load_dataset(\"stabilityai/word-list\", data_files=\"list.txt\", use_auth_token=True)\n",
    "#word_list = word_list_dataset[\"train\"]['text']\n",
    "WORD_LIST = 'https://raw.githubusercontent.com/coffee-and-fun/google-profanity-words/main/data/list.txt'\n",
    "import requests\n",
    "word_list = [word for word in requests.get(WORD_LIST).text.split('\\n') if word and not word.isspace()]\n",
    "\n",
    "def infer(prompt, samples, steps, scale, seed):\n",
    "    for filter in word_list:\n",
    "        if re.search(rf\"\\b{filter}\\b\", prompt):\n",
    "            raise gr.Error(f\"\"\"Unsafe content found. Please try again with different prompts. \n",
    "            filter: {filter}\n",
    "            prompt: {prompt}\n",
    "\"\"\")\n",
    "\n",
    "    generator = torch.Generator(device=device).manual_seed(seed)\n",
    "    \n",
    "    with torch.autocast(pipe.device.type):\n",
    "        yield from pipe.generate(\n",
    "            [prompt] * samples,\n",
    "            num_inference_steps=steps,\n",
    "            guidance_scale=scale,\n",
    "            generator=generator,\n",
    "        )\n",
    "\n",
    "\n",
    "def replace_unsafe_images(output):\n",
    "    images = []\n",
    "    safe_image = Image.open(r\"unsafe.png\")\n",
    "    for image, is_unsafe in zip(output.images, output.nsfw_content_detected):\n",
    "        if is_unsafe:\n",
    "            images.append(safe_image)\n",
    "        else:\n",
    "            images.append(pipe.numpy_to_pil(image)[0])\n",
    "    return images\n",
    "\n",
    "\n",
    "approximate_decoder = ApproximateDecoder.for_pipeline(pipe)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "examples = [\n",
    "    [\n",
    "        'A high tech solarpunk utopia in the Amazon rainforest',\n",
    "        2,\n",
    "        20,\n",
    "        7.5,\n",
    "        1024,\n",
    "    ],\n",
    "    [\n",
    "        'A pikachu fine dining with a view to the Eiffel Tower',\n",
    "        2,\n",
    "        20,\n",
    "        7,\n",
    "        1024,\n",
    "    ],\n",
    "    [\n",
    "        'A mecha robot in a favela in expressionist style',\n",
    "        2,\n",
    "        20,\n",
    "        7,\n",
    "        1024,\n",
    "    ],\n",
    "    [\n",
    "        'an insect robot preparing a delicious meal',\n",
    "        2,\n",
    "        20,\n",
    "        7,\n",
    "        1024,\n",
    "    ],\n",
    "    [\n",
    "        \"A small cabin on top of a snowy mountain in the style of Disney, artstation\",\n",
    "        2,\n",
    "        20,\n",
    "        7,\n",
    "        1024,\n",
    "    ],\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "import urllib.parse\n",
    "icon_src = ('<svg viewBox=\"0 0 115 115\" fill=\"none\" xmlns=\"http://www.w3.org/2000/svg\">'\n",
    "    '<path fill=\"#fff\" d=\"M0 0h23v23H0zm0 69h23v23H0z\"/><path fill=\"#AEAEAE\" d=\"M23 0h23v23H23zm0 69h23v23H23z\"/>'\n",
    "    '<path fill=\"#fff\" d=\"M46 0h23v23H46zm0 69h23v23H46z\"/><path fill=\"#000\" d=\"M69 0h23v23H69zm0 69h23v23H69z\"/>'\n",
    "    '<path fill=\"#D9D9D9\" d=\"M92 0h23v23H92z\"/><path fill=\"#AEAEAE\" d=\"M92 69h23v23H92z\"/><path fill=\"#fff\" d=\"M115 46h23v23h-23zm0 69h23v23h-23z\"/>'\n",
    "    '<path fill=\"#D9D9D9\" d=\"M115 69h23v23h-23z\"/><path fill=\"#AEAEAE\" d=\"M92 46h23v23H92zm0 69h23v23H92z\"/>'\n",
    "    '<path fill=\"#fff\" d=\"M92 69h23v23H92zM69 46h23v23H69zm0 69h23v23H69z\"/><path fill=\"#D9D9D9\" d=\"M69 69h23v23H69z\"/>'\n",
    "    '<path fill=\"#000\" d=\"M46 46h23v23H46zm0 69h23v23H46zm0-46h23v23H46z\"/><path fill=\"#D9D9D9\" d=\"M23 46h23v23H23z\"/>'\n",
    "    '<path fill=\"#AEAEAE\" d=\"M23 115h23v23H23z\"/><path fill=\"#000\" d=\"M23 69h23v23H23z\"/></svg>')\n",
    "icon = \"data:image+svg/xml,\" + urllib.parse.quote(icon_src, safe=\" \\\"=\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "679f1dc87b194c078d7a4e44cf44ea6e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HBox(children=(Text(value='', placeholder='Enter your prompt'), Button(description='Generate im…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import secrets\n",
    "import ipywidgets\n",
    "\n",
    "PREVIEW_ZOOM = 3\n",
    "DEFAULT_SIZE = 512\n",
    "\n",
    "_latent_size = DEFAULT_SIZE >> 3\n",
    "\n",
    "def _preview_widget():\n",
    "    widget = ipywidgets.Image(\n",
    "        format='jpeg',\n",
    "        width=_latent_size,\n",
    "        height=_latent_size,\n",
    "    )\n",
    "    preview_size = f\"{PREVIEW_ZOOM * _latent_size}px\"\n",
    "    widget.layout.width = preview_size\n",
    "    widget.layout.height = preview_size\n",
    "    widget.layout.object_position = 'center center'\n",
    "    return widget\n",
    "    \n",
    "def on_submit(event=None):\n",
    "    gallery_container.children = [_preview_widget() for i in range(samples.value)]\n",
    "    progress.value = progress.max\n",
    "    for result in infer(text.value, samples.value, steps.value, scale.value, seed.value):\n",
    "        if isinstance(result, PipelineIntermediateState):\n",
    "            progress.value = result.timestep\n",
    "            for widget, latents in zip(gallery_container.children, result.latents):\n",
    "                widget.value = jpeg_bytes(approximate_decoder(latents))\n",
    "        else:\n",
    "            progress.value = progress.min\n",
    "            widgets = []\n",
    "            for image in replace_unsafe_images(result):\n",
    "                widgets.append(ipywidgets.Image(\n",
    "                    value=jpeg_bytes(image), \n",
    "                    format='jpeg', \n",
    "                    width=512, height=512\n",
    "                ))\n",
    "            gallery_container.children = widgets\n",
    "            \n",
    "\n",
    "text = ipywidgets.Text(\n",
    "    # description=\"Enter your prompt\",\n",
    "    placeholder=\"Enter your prompt\"\n",
    ")\n",
    "btn = ipywidgets.Button(\n",
    "    description=\"Generate image\"\n",
    ")\n",
    "btn.on_click(on_submit)\n",
    "text.on_submit(on_submit)\n",
    "gallery_container = ipywidgets.Box(layout=ipywidgets.Layout(\n",
    "    flex_flow=\"row wrap\",\n",
    "    justify_content=\"space-around\",\n",
    "    align_items=\"center\",\n",
    "    align_content=\"space-around\",    \n",
    "))\n",
    "\n",
    "MAX_SEED = 1 << 31 - 1\n",
    "samples = ipywidgets.IntSlider(2, 1, 4, description=\"Images\")\n",
    "steps = ipywidgets.IntSlider(16, 1, 50, description=\"Steps\")\n",
    "scale = ipywidgets.FloatSlider(7.5, min=0, max=50, step=0.1, description=\"Guidance Scale\")\n",
    "seed = ipywidgets.BoundedIntText(secrets.randbelow(MAX_SEED), 0, MAX_SEED, description=\"Seed\")\n",
    "progress = ipywidgets.IntProgress(value=0, min=0, max=pipe.scheduler.num_train_timesteps)\n",
    "progress.layout.width = \"100%\"\n",
    "\n",
    "advanced_options = ipywidgets.Accordion(children=[ipywidgets.VBox([\n",
    "    samples,\n",
    "    steps,\n",
    "    scale,\n",
    "    seed\n",
    "])])\n",
    "advanced_options.set_title(0, \"Advanced Options\")\n",
    "\n",
    "form = ipywidgets.VBox([\n",
    "    ipywidgets.HBox([text, btn]),\n",
    "    gallery_container,\n",
    "    progress,\n",
    "    advanced_options\n",
    "])\n",
    "display(form)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model by [CompVis][CompVis] and [Stability AI][Stability AI].\n",
    "\n",
    "### LICENSE\n",
    "\n",
    "The model is licensed with a [CreativeML Open RAIL-M][license] license.\n",
    "The authors claim no rights on the outputs you generate,\n",
    "you are free to use them and are accountable for their use which must not go against the provisions set in this license.\n",
    "The license forbids you from sharing any content that violates any laws, produce any harm to a person,\n",
    "disseminate any personal information that would be meant for harm,\n",
    "spread misinformation and target vulnerable groups.\n",
    "For the full list of restrictions please [read the license][license]\n",
    "\n",
    "### Biases and content acknowledgment\n",
    "\n",
    "Despite how impressive being able to turn text into image is,\n",
    "beware to the fact that this model may output content that reinforces or exacerbates societal biases,\n",
    "as well as realistic faces, pornography and violence.\n",
    "The model was trained on the [LAION-5B dataset][laion-5b],\n",
    "which scraped non-curated image-text-pairs from the internet (the exception being the removal of illegal content)\n",
    "and is meant for research purposes.\n",
    "You can read more in the [model card][card].\n",
    "\n",
    "[CompVis]: https://huggingface.co/CompVis\n",
    "[Stability AI]: https://huggingface.co/stabilityai\n",
    "[license]: https://huggingface.co/spaces/CompVis/stable-diffusion-license\n",
    "[laion-5b]: https://laion.ai/blog/laion-5b/\n",
    "[card]: https://huggingface.co/CompVis/stable-diffusion-v1-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "diffusers-demo",
   "language": "python",
   "name": "diffusers-demo"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}