quick_samples.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98b83c8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import StableDiffusionPipeline, UNet2DConditionModel, StableDiffusionXLPipeline\n",
    "import torch\n",
    "torch.set_grad_enabled(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "192c556e",
   "metadata": {},
   "outputs": [],
   "source": [
    "dpo_unet = UNet2DConditionModel.from_pretrained(\n",
    "                            #  'mhdang/dpo-sd1.5-text2image-v1',\n",
    "                            'mhdang/dpo-sdxl-text2image-v1',\n",
    "                            # alternatively use local ckptdir (*/checkpoint-n/)\n",
    "                            subfolder='unet',\n",
    "                            torch_dtype=torch.float16\n",
    ").to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e06ccfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pretrained_model_name = \"CompVis/stable-diffusion-v1-4\"\n",
    "# pretrained_model_name = \"runwayml/stable-diffusion-v1-5\"\n",
    "pretrained_model_name = \"stabilityai/stable-diffusion-xl-base-1.0\"\n",
    "gs = (5 if 'stable-diffusion-xl' in pretrained_model_name else 7.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e28d893e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if 'stable-diffusion-xl' in pretrained_model_name:\n",
    "    pipe = StableDiffusionXLPipeline.from_pretrained(\n",
    "        pretrained_model_name, torch_dtype=torch.float16,\n",
    "        variant=\"fp16\", use_safetensors=True\n",
    "    ).to(\"cuda\")\n",
    "else:\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(pretrained_model_name,\n",
    "                                                   torch_dtype=torch.float16)\n",
    "pipe = pipe.to('cuda')\n",
    "pipe.safety_checker = None # Trigger-happy, blacks out >50% of \"robot tiger\"\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2150c454",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Can do clip_utils, aes_utils, hps_utils\n",
    "from utils.pickscore_utils import Selector\n",
    "# Score generations automatically w/ reward model\n",
    "ps_selector = Selector('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe5edb88",
   "metadata": {},
   "outputs": [],
   "source": [
    "unets = [pipe.unet, dpo_unet]\n",
    "names = [\"Orig. SDXL\", \"DPO SDXL\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64ef0616",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen(prompt, seed=0, run_baseline=True):\n",
    "    ims = []\n",
    "    generator = torch.Generator(device='cuda')\n",
    "    for unet_i in ([0, 1] if run_baseline else [1]):\n",
    "        print(f\"Prompt: {prompt}\\nSeed: {seed}\\n{names[unet_i]}\")\n",
    "        pipe.unet = unets[unet_i]\n",
    "        generator = generator.manual_seed(seed)\n",
    "        \n",
    "        im = pipe(prompt=prompt, generator=generator, guidance_scale=gs).images[0]\n",
    "        display(im)\n",
    "        ims.append(im)\n",
    "    return ims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0bb84b9",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "example_prompts = [\n",
    "    \"A pile of sand swirling in the wind forming the shape of a dancer\",\n",
    "    \"A giant dinosaur frozen into a glacier and recently discovered by scientists, cinematic still\",\n",
    "    \"a smiling beautiful sorceress with long dark hair and closed eyes wearing a dark top surrounded by glowing fire sparks at night, magical light fog, deep focus+closeup, hyper-realistic, volumetric lighting, dramatic lighting, beautiful composition, intricate details, instagram, trending, photograph, film grain and noise, 8K, cinematic, post-production\",\n",
    "    \"A purple raven flying over big sur, light fog, deep focus+closeup, hyper-realistic, volumetric lighting, dramatic lighting, beautiful composition, intricate details, instagram, trending, photograph, film grain and noise, 8K, cinematic, post-production\",\n",
    "    \"a smiling beautiful sorceress wearing a modest high necked blue suit surrounded by swirling rainbow aurora, hyper-realistic, cinematic, post-production\",\n",
    "    \"Anthro humanoid turtle skydiving wearing goggles, gopro footage\",\n",
    "    \"A man in a suit surfing in a river\",\n",
    "    \"photo of a zebra dressed suit and tie sitting at a table in a bar with a bar stools, award winning photography\",\n",
    "    \"A typhoon in a tea cup, digital render\",\n",
    "    \"A cute puppy leading a session of the United Nations, newspaper photography\",\n",
    "    \"Worm eye view of rocketship\",\n",
    "    \"Glass spheres in the desert, refraction render\",\n",
    "    \"anthropmorphic coffee bean drinking coffee\",\n",
    "    \"A baby kangaroo in a trenchcoat\",\n",
    "    \"A towering hurricane of rainbow colors towering over a city, cinematic digital art\",\n",
    "    \"A redwood tree rising up out of the ocean\",\n",
    "]\n",
    "\n",
    "\n",
    "for p in example_prompts:\n",
    "    ims = gen(p) # could save these if desired    \n",
    "    scores = ps_selector.score(ims, p)\n",
    "    print(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c94d034",
   "metadata": {},
   "outputs": [],
   "source": [
    "# to get partiprompts captions\n",
    "from datasets import load_dataset\n",
    "dataset = load_dataset(\"nateraw/parti-prompts\")\n",
    "print(dataset['train']['Prompt'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}