CLIPDraw_3DStrokes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e7DCzf-EHzL_"
   },
   "source": [
    "# CLIPDraw with 3d strokes\n",
    "\n",
    "Synthesize (3d) drawings to match a text prompt.\n",
    "\n",
    "This places the output in created directories tmp/ and output/, but doesn't display the results inside the notebook itself.\n",
    "\n",
    "To make a rotating view, after training for as long as you like, run the cell under the Rotating View header\n",
    "\n",
    "This version is set up to use a different prompt for different angles, though you can set all the prompts identical to avoid this effect."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports and setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To use in Colab, probably need this...\n",
    "\n",
    "#!pip install clip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hjt9T3ARukAg",
    "outputId": "e0f9ed7c-3e31-41ee-cf7d-629f0f6f553d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mkdir: cannot create directory ‘tmp’: File exists\n",
      "mkdir: cannot create directory ‘output’: File exists\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import io\n",
    "import numpy as np\n",
    "import glob\n",
    "import subprocess\n",
    "from math import *\n",
    "import clip\n",
    "\n",
    "from PIL import Image\n",
    "import random\n",
    "#import ttools.modules\n",
    "import argparse\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torchvision import transforms\n",
    "import torchvision\n",
    "\n",
    "!mkdir tmp\n",
    "!mkdir output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Z-Wt7UjTi8Le",
    "outputId": "f626931b-5130-4c7d-8bef-18a77b40e28e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2343, 512]) torch.float16\n"
     ]
    }
   ],
   "source": [
    "#@title Load CLIP {vertical-output: true}\n",
    "\n",
    "# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'\n",
    "\n",
    "# Load the model\n",
    "device = torch.device('cuda')\n",
    "model, preprocess = clip.load('ViT-B/32', device, jit=False)\n",
    "\n",
    "nouns = \"aardvark abyssinian accelerator accordion account accountant acknowledgment acoustic acrylic act action active activity actor actress adapter addition address adjustment adult advantage advertisement advice afghanistan africa aftermath afternoon aftershave afterthought age agenda agreement air airbus airmail airplane airport airship alarm albatross alcohol algebra algeria alibi alley alligator alloy almanac alphabet alto aluminium aluminum ambulance america amount amusement anatomy anethesiologist anger angle angora animal anime ankle answer ant antarctica anteater antelope anthony anthropology apartment apology apparatus apparel appeal appendix apple appliance approval april aquarius arch archaeology archeology archer architecture area argentina argument aries arithmetic arm armadillo armchair armenian army arrow art ash ashtray asia asparagus asphalt asterisk astronomy athlete atm atom attack attempt attention attic attraction august aunt australia australian author authorisation authority authorization avenue babies baboon baby back backbone bacon badge badger bag bagel bagpipe bail bait baker bakery balance balinese ball balloon bamboo banana band bandana bangladesh bangle banjo bank bankbook banker bar barbara barber barge baritone barometer base baseball basement basin basket basketball bass bassoon bat bath bathroom bathtub battery battle bay beach bead beam bean bear beard beast beat beautician beauty beaver bed bedroom bee beech beef beer beet beetle beggar beginner begonia behavior belgian belief believe bell belt bench bengal beret berry bestseller betty bibliography bicycle bike bill billboard biology biplane birch bird birth birthday bit bite black bladder blade blanket blinker blizzard block blood blouse blow blowgun blue board boat bobcat body bolt bomb bomber bone bongo bonsai book bookcase booklet boot border botany bottle bottom boundary bow bowl bowling box boy bra brace bracket brain brake branch brand brandy brass brazil bread break breakfast breath brian brick bridge british broccoli brochure broker bronze brother brother-in-law brow brown brush bubble bucket budget buffer buffet bugle building bulb bull bulldozer bumper bun burglar burma burn burst bus bush business butane butcher butter button buzzard cabbage cabinet cable cactus cafe cake calculator calculus calendar calf call camel camera camp can canada canadian cancer candle cannon canoe canvas cap capital cappelletti capricorn captain caption car caravan carbon card cardboard cardigan care carnation carol carp carpenter carriage carrot cart cartoon case cast castanet cat catamaran caterpillar cathedral catsup cattle cauliflower cause caution cave c-clamp cd ceiling celery celeste cell cellar cello celsius cement cemetery cent centimeter century ceramic cereal certification chain chair chalk chance change channel character chard charles chauffeur check cheek cheese cheetah chef chemistry cheque cherries cherry chess chest chick chicken chicory chief child children chill chime chimpanzee chin china chinese chive chocolate chord christmas christopher chronometer church cicada cinema circle circulation cirrus citizenship city clam clarinet class claus clave clef clerk click client climb clipper cloakroom clock close closet cloth cloud cloudy clover club clutch coach coal coast coat cobweb cockroach cocktail cocoa cod coffee coil coin coke cold collar college collision colombia colon colony color colt column columnist comb comfort comic comma command commission committee community company comparison competition competitor composer composition computer condition condor cone confirmation conga congo conifer connection consonant continent control cook cooking copper copy copyright cord cork cormorant corn cornet correspondent cost cotton couch cougar cough country course court cousin cover cow cowbell crab crack cracker craftsman crate crawdad crayfish crayon cream creator creature credit creditor creek crib cricket crime criminal crocodile crocus croissant crook crop cross crow crowd crown crush cry cub cuban cucumber cultivator cup cupboard cupcake curler currency current curtain curve cushion custard customer cut cuticle cycle cyclone cylinder cymbal dad daffodil dahlia daisy damage dance dancer danger daniel dash dashboard database date daughter david day dead deadline deal death deborah debt debtor decade december decimal decision decrease dedication deer defense deficit degree delete delivery den denim dentist deodorant department deposit description desert design desire desk dessert destruction detail detective development dew diamond diaphragm dibble dictionary dietician difference digestion digger digital dill dime dimple dinghy dinner dinosaur diploma dipstick direction dirt disadvantage discovery discussion disease disgust dish distance distribution distributor diving division divorced dock doctor dog dogsled doll dollar dolphin domain donald donkey donna door dorothy double doubt downtown dragon dragonfly drain drake drama draw drawbridge drawer dream dredger dress dresser dressing drill drink drive driver driving drizzle drop drug drum dry dryer duck duckling dugout dungeon dust eagle ear earth earthquake ease east edge edger editor editorial education edward eel effect egg eggnog eggplant egypt eight elbow element elephant elizabeth ellipse emery employee employer encyclopedia end enemy energy engine engineer engineering english enquiry entrance environment epoch epoxy equinox equipment era error estimate ethernet ethiopia euphonium europe evening event examination example exchange exclamation exhaust ex-husband existence expansion experience expert explanation ex-wife eye eyebrow eyelash eyeliner face facilities fact factory fahrenheit fairies fall family fan fang farm farmer fat father father-in-law faucet fear feast feather feature february fedelini feedback feeling feet felony female fender ferry ferryboat fertilizer fiber fiberglass fibre fiction field fifth fight fighter file find fine finger fir fire fired fireman fireplace firewall fish fisherman flag flame flare flat flavor flax flesh flight flock flood floor flower flugelhorn flute fly foam fog fold font food foot football footnote force forecast forehead forest forgery fork form format fortnight foundation fountain fowl fox foxglove fragrance frame france freckle freeze freezer freighter french freon friction friday fridge friend frog front frost frown fruit fuel fur furniture galley gallon game gander garage garden garlic gas gasoline gate gateway gauge gazelle gear gearshift geese gemini gender geography geology geometry george geranium german germany ghana ghost giant giraffe girdle girl gladiolus glass glider gliding glockenspiel glove glue goal goat gold goldfish golf gondola gong good-bye goose gore-tex gorilla gosling government governor grade grain gram granddaughter grandfather grandmother grandson grape graphic grass grasshopper gray grease great-grandfather great-grandmother greece greek green grenade grey grill grip ground group grouse growth guarantee guatemalan guide guilty guitar gum gun gym gymnast hacksaw hail hair haircut half-brother half-sister halibut hall hallway hamburger hammer hamster hand handball handicap handle handsaw harbor hardboard hardcover hardhat hardware harmonica harmony harp hat hate hawk head headlight headline health hearing heart heat heaven hedge height helen helicopter helium hell helmet help hemp hen heron herring hexagon hill himalayan hip hippopotamus history hobbies hockey hoe hole holiday home honey hood hook hope horn horse hose hospital hot hour hourglass house hovercraft hub hubcap humidity humor hurricane hyacinth hydrant hydrofoil hydrogen hyena hygienic ice icebreaker icicle icon idea ikebana illegal imprisonment improvement impulse inch income increase index india indonesia industry ink innocent input insect instruction instrument insulation insurance interactive interest internet interviewer intestine invention inventory invoice iran iraq iris iron island israel italian italy jacket jaguar jail jam james january japan japanese jar jasmine jason jaw jeans jeep jeff jelly jellyfish jennifer jet jewel jogging john join joke joseph journey judge judo juice july jumbo jump jumper june jury justice jute kale kamikaze kangaroo karate karen kayak kendo kenneth kenya ketchup kettle kettledrum kevin key keyboard keyboarding kick kidney kilogram kilometer kimberly kiss kitchen kite kitten kitty knee knickers knife knight knot knowledge kohlrabi korean laborer lace ladybug lake lamb lamp lan land landmine language larch lasagna latency latex lathe laugh laundry laura law lawyer layer lead leaf learning leather leek leg legal lemonade lentil leo leopard letter lettuce level libra library license lier lift light lightning lilac lily limit linda line linen link lion lip lipstick liquid liquor lisa list literature litter liver lizard llama loaf loan lobster lock locket locust look loss lotion love low lumber lunch lunchroom lung lunge lute luttuce lycra lynx lyocell lyre lyric macaroni machine macrame magazine magic magician maid mail mailbox mailman makeup malaysia male mall mallet man manager mandolin manicure manx map maple maraca marble march margaret margin maria marimba mark mark market married mary mascara mask mass match math mattock may mayonnaise meal measure meat mechanic medicine meeting melody memory men menu mercury message metal meteorology meter methane mexican mexico mice michael michelle microwave middle mile milk milkshake millennium millimeter millisecond mimosa mind mine minibus mini-skirt minister mint minute mirror missile mist mistake mitten moat modem mole mom monday money monkey month moon morning morocco mosque mosquito mother mother-in-law motion motorboat motorcycle mountain mouse moustache mouth move multi-hop multimedia muscle museum music musician mustard myanmar nail name nancy napkin narcissus nation neck need needle neon nepal nephew nerve nest net network news newsprint newsstand nic nickel niece nigeria night nitrogen node noise noodle north north america north korea norwegian nose note notebook notify novel november number numeric nurse nut nylon oak oatmeal objective oboe observation occupation ocean ocelot octagon octave october octopus odometer offence offer office oil okra olive onion open opera operation ophthalmologist opinion option orange orchestra orchid order organ organisation organization ornament ostrich otter ounce output outrigger oval oven overcoat owl owner ox oxygen oyster package packet page pail pain paint pair pajama pakistan palm pamphlet pan pancake pancreas panda pansy panther panties pantry pants panty pantyhose paper paperback parade parallelogram parcel parent parentheses park parrot parsnip part particle partner partridge party passbook passenger passive pasta paste pastor pastry patch path patient patio patricia paul payment pea peace peak peanut pear pedestrian pediatrician peen peer-to-peer pelican pen penalty pencil pendulum pentagon peony pepper perch perfume period periodical peripheral permission persian person peru pest pet pharmacist pheasant philippines philosophy phone physician piano piccolo pickle picture pie pig pigeon pike pillow pilot pimple pin pine ping pink pint pipe pisces pizza place plain plane planet plant plantation plaster plasterboard plastic plate platinum play playground playroom pleasure plier plot plough plow plywood pocket poet point poison poland police policeman polish politician pollution polo polyester pond popcorn poppy population porch porcupine port porter position possibility postage postbox pot potato poultry pound powder power precipitation preface prepared pressure price priest print printer prison probation process processing produce product production professor profit promotion propane property prose prosecution protest protocol pruner psychiatrist psychology ptarmigan puffin pull puma pump pumpkin punch punishment puppy purchase purple purpose push pvc pyjama pyramid quail quality quart quarter quartz queen question quicksand quiet quill quilt quince quit quiver quotation rabbi rabbit racing radar radiator radio radish raft rail railway rain rainbow raincoat rainstorm rake ramie random range rat rate raven ravioli ray rayon reaction reading reason receipt recess record recorder rectangle red reduction refrigerator refund regret reindeer relation relative religion relish reminder repair replace report representative request resolution respect responsibility rest restaurant result retailer revolve revolver reward rhinoceros rhythm rice richard riddle rifle ring rise risk river riverbed road roadway roast robert robin rock rocket rod roll romania romanian ronald roof room rooster root rose rotate route router rowboat rub rubber rugby rule run russia russian rutabaga ruth sack sagittarius sail sailboat sailor salad salary sale salesman salmon salt sampan samurai sand sandra sandwich santa sarah sardine satin saturday sauce saudi arabia sausage save saw saxophone scale scallion scanner scarecrow scarf scene scent schedule school science scissors scooter scorpio scorpion scraper screen screw screwdriver sea seagull seal seaplane search seashore season seat second secretary secure security seed seeder segment select selection self semicircle semicolon sense sentence separated september servant server session sex shade shadow shake shallot shame shampoo shape share shark sharon shears sheep sheet shelf shell shield shingle ship shirt shock shoe shoemaker shop shorts shoulder shovel show shrimp shrine siamese siberian side sideboard sidecar sidewalk sign signature silica silk silver sing singer single sink sister sister-in-law size skate skiing skill skin skirt sky slash slave sled sleep sleet slice slime slip slipper slope smash smell smile smoke snail snake sneeze snow snowboarding snowflake snowman snowplow snowstorm soap soccer society sociology sock soda sofa softball softdrink software soil soldier son song soprano sort sound soup sousaphone south africa south america south korea soy soybean space spade spaghetti spain spandex spark sparrow spear specialist speedboat sphere sphynx spider spike spinach spleen sponge spoon spot spring sprout spruce spy square squash squid squirrel stage staircase stamp star start starter state statement station statistic steam steel stem step step-aunt step-brother stepdaughter step-daughter step-father step-grandfather step-grandmother stepmother step-mother step-sister stepson step-son step-uncle steven stew stick stinger stitch stock stocking stomach stone stool stop stopsign stopwatch store storm story stove stranger straw stream street streetcar stretch string structure study sturgeon submarine substance subway success sudan suede sugar suggestion suit summer sun sunday sundial sunflower sunshine supermarket supply support surfboard surgeon surname surprise susan sushi swallow swamp swan sweater sweatshirt sweatshop swedish sweets swim swimming swing swiss switch sword swordfish sycamore syria syrup system table tablecloth tabletop tachometer tadpole tail tailor taiwan talk tank tanker tanzania target taste taurus tax taxi taxicab tea teacher teaching team technician teeth television teller temper temperature temple tempo tendency tennis tenor tent territory test text textbook texture thailand theater theory thermometer thing thistle thomas thought thread thrill throat throne thumb thunder thunderstorm thursday ticket tie tiger tights tile timbale time timer timpani tin tip tire titanium title toad toast toe toenail toilet tomato tom-tom ton tongue tooth toothbrush toothpaste top tornado tortellini tortoise touch tower town toy tractor trade traffic trail train tramp transaction transmission transport trapezoid tray treatment tree trial triangle trick trigonometry trip trombone trouble trousers trout trowel truck trumpet trunk t-shirt tsunami tub tuba tuesday tugboat tulip tuna tune turkey turkey turkish turn turnip turnover turret turtle tv twig twilight twine twist typhoon tyvek uganda ukraine ukrainian umbrella uncle underclothes underpants undershirt underwear unit united kingdom unshielded use utensil uzbekistan vacation vacuum valley value van vase vault vegetable vegetarian veil vein velvet venezuela venezuelan verdict vermicelli verse vessel vest veterinarian vibraphone vietnam view vinyl viola violet violin virgo viscose vise vision visitor voice volcano volleyball voyage vulture waiter waitress walk wall wallaby wallet walrus war warm wash washer wasp waste watch watchmaker water waterfall wave wax way wealth weapon weasel weather wedge wednesday weed weeder week weight whale wheel whip whiskey whistle white wholesaler whorl wilderness william willow wind windchime window windscreen windshield wine wing winter wire wish witch withdrawal witness wolf woman women wood wool woolen word work workshop worm wound wrecker wren wrench wrinkle wrist writer xylophone yacht yak yam yard yarn year yellow yew yogurt yoke yugoslavian zebra zephyr zinc zipper zone zoo zoology\"\n",
    "nouns = nouns.split(\" \")\n",
    "noun_prompts = [\"a drawing of a \" + x for x in nouns]\n",
    "\n",
    "# Calculate features\n",
    "with torch.no_grad():\n",
    "    nouns_features = model.encode_text(torch.cat([clip.tokenize(noun_prompts).to(device)]))\n",
    "print(nouns_features.shape, nouns_features.dtype)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Differentiable 3d stroke renderer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_blobs(cam_pos, cam_targ, xyz, rgba, width, density=2):\n",
    "    xx,yy = np.meshgrid(np.arange(canvas_width), np.arange(canvas_height))\n",
    "    canv_coords = np.concatenate([yy[:,:,np.newaxis], xx[:,:,np.newaxis]], axis=2)\n",
    "    canv_coords = torch.FloatTensor(canv_coords).to(device)\n",
    "    \n",
    "    up = torch.FloatTensor([0,1,0]).to(device)\n",
    "    \n",
    "    cam_u = F.normalize(cam_targ-cam_pos, dim=0)\n",
    "    cam_v = torch.cross(cam_u, up)\n",
    "    cam_w = torch.cross(cam_u, cam_v)\n",
    "    \n",
    "    pos = xyz - cam_pos.unsqueeze(0)\n",
    "    pos_u = torch.sum(pos*cam_u.unsqueeze(0), 1)\n",
    "    pos_v = torch.sum(pos*cam_v.unsqueeze(0), 1)\n",
    "    pos_w = torch.sum(pos*cam_w.unsqueeze(0), 1)\n",
    "    \n",
    "    pos_x = canvas_width * pos_w / (1e-4 + pos_u) + canvas_width // 2\n",
    "    pos_y = canvas_width * pos_v / (1e-4 + pos_u) + canvas_height // 2\n",
    "    pos_xy = torch.cat([pos_x.unsqueeze(1), pos_y.unsqueeze(1)], 1)\n",
    "    modified_w = canvas_width * width / (1e-4 + pos_u)\n",
    "    \n",
    "    blobs_rgb = []\n",
    "    blobs_alpha = []\n",
    "    blob_depths = []\n",
    "    blob_offsets = []\n",
    "    \n",
    "    for i in range(xyz.shape[0]-1):\n",
    "        dist = torch.sqrt(torch.sum( (pos_xy[i] - pos_xy[i+1])**2 ))\n",
    "        wbar = (modified_w[i] + modified_w[i+1])*0.5\n",
    "        steps = int((density * dist / (1 + wbar)).cpu().detach().item()) + 1\n",
    "        if steps<2:\n",
    "            steps = 2\n",
    "        if steps>15:\n",
    "            steps = 15\n",
    "        \n",
    "        for j in range(steps):\n",
    "            s = j/(steps-1)\n",
    "            \n",
    "            p = pos_xy[i]*(1-s) + pos_xy[i+1]*s\n",
    "            c = rgba[i]*(1-s) + rgba[i+1]*s\n",
    "            w = modified_w[i]*(1-s) + modified_w[i+1]*s\n",
    "            d = pos_u[i]*(1-s) + pos_u[i+1]*s\n",
    "            \n",
    "            if d>0.05: # Culling                \n",
    "                bounds = 2*int(2*w.cpu().detach().item())+2\n",
    "                offset = (p.cpu().detach().numpy()-bounds//2).astype(np.int32)\n",
    "                if offset[0]<0:\n",
    "                    offset[0]=0\n",
    "                if offset[1]<0:\n",
    "                    offset[1]=0\n",
    "                \n",
    "                cc = canv_coords[offset[0]:offset[0]+bounds, offset[1]:offset[1]+bounds]\n",
    "                \n",
    "                r2 = torch.sum( (cc-p.unsqueeze(0).unsqueeze(1))**2, 2)/w**2\n",
    "\n",
    "                r2 = r2.unsqueeze(2)\n",
    "                \n",
    "                blob_a = c[3:4].unsqueeze(0).unsqueeze(1) * torch.exp(-r2**2)\n",
    "                \n",
    "                blob_offsets.append(offset)\n",
    "                blob_depths.append(d.cpu().detach().item())                \n",
    "                blobs_rgb.append(c[:3].unsqueeze(1))\n",
    "                blobs_alpha.append(blob_a)\n",
    "    \n",
    "    return blob_depths, blobs_rgb, blobs_alpha, blob_offsets\n",
    "\n",
    "def bounded_composition(canvas, offset, rgb, alpha):\n",
    "    # Need to split out the affected square, then concatenate everything back\n",
    "    mod_part = canvas[offset[0]:offset[0]+alpha.shape[0], offset[1]:offset[1]+alpha.shape[1]]\n",
    "    \n",
    "    mod_part = (1-alpha)*mod_part + alpha * rgb.unsqueeze(0).unsqueeze(1)\n",
    "    \n",
    "    pre_x = canvas[:offset[0], :, :]\n",
    "    \n",
    "    top_bit = canvas[offset[0]:offset[0]+alpha.shape[0], :offset[1], :]\n",
    "    bottom_bit = canvas[offset[0]:offset[0]+alpha.shape[0], offset[1]+alpha.shape[1]:, :]\n",
    "    \n",
    "    post_x = canvas[offset[0]+alpha.shape[0]:, :, :]\n",
    "    \n",
    "    central_column = torch.cat([top_bit, mod_part, bottom_bit], 1)\n",
    "    \n",
    "    canvas = torch.cat([pre_x, central_column, post_x], 0)\n",
    "    \n",
    "    return canvas\n",
    "    \n",
    "def compose_blobs(depths, rgbs, alphas, offsets):\n",
    "    canvas = torch.ones(canvas_width, canvas_height, 3).to(device)\n",
    "    \n",
    "    idx = np.argsort(-np.array(depths))\n",
    "    rgbs = torch.cat(rgbs, 1)   \n",
    "    \n",
    "    for i in range(rgbs.shape[1]):\n",
    "        canvas = bounded_composition(canvas, offsets[idx[i]], rgbs[:,idx[i]], alphas[idx[i]])\n",
    "        \n",
    "    return canvas\n",
    "\n",
    "def render(cam_pos, cam_targ, strokes):\n",
    "    depths = []\n",
    "    rgbs = []\n",
    "    alphas = []\n",
    "    offsets = []\n",
    "    for s in strokes:\n",
    "        bd,br,ba,bo = get_blobs(cam_pos, cam_targ, s[0], s[1], s[2])\n",
    "        depths = depths + bd\n",
    "        rgbs = rgbs + br\n",
    "        alphas = alphas + ba\n",
    "        offsets = offsets + bo\n",
    "\n",
    "    img = compose_blobs(depths, rgbs, alphas, offsets).view(canvas_width, canvas_height, 3)\n",
    "    \n",
    "    del depths\n",
    "    del rgbs\n",
    "    del alphas\n",
    "    del offsets\n",
    "    \n",
    "    return img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Enter prompts here!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "4XIVMSJuWgxG"
   },
   "outputs": [],
   "source": [
    "#@title Curve Optimizer {vertical-output: true}\n",
    "\n",
    "prompt1 = \"The silhouette of a rabbit\"\n",
    "text_input1 = clip.tokenize(prompt1).to(device)\n",
    "\n",
    "prompt2 = \"A vase containing flowers\"\n",
    "text_input2 = clip.tokenize(prompt2).to(device)\n",
    "\n",
    "prompt3 = \"A waluigi\"\n",
    "text_input3 = clip.tokenize(prompt3).to(device)\n",
    "\n",
    "# Calculate features\n",
    "with torch.no_grad():\n",
    "    text_features1 = model.encode_text(text_input1)\n",
    "    text_features2 = model.encode_text(text_input2)\n",
    "    text_features3 = model.encode_text(text_input3)\n",
    "\n",
    "gamma = 1.0\n",
    "\n",
    "args = lambda: None\n",
    "args.num_paths = 200\n",
    "args.max_width = 0.5\n",
    "args.num_iter = 500\n",
    "args.use_blob = False\n",
    "args.use_lpips_loss = False\n",
    "\n",
    "# Use GPU if available\n",
    "device = torch.device('cuda')\n",
    "\n",
    "canvas_width, canvas_height = 224, 224\n",
    "num_paths = args.num_paths\n",
    "max_width = args.max_width\n",
    "\n",
    "augment_trans = transforms.Compose([\n",
    "    transforms.RandomPerspective(fill=1, p=1),\n",
    "    transforms.RandomResizedCrop(224, scale=(0.7,0.9)),\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize strokes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Strokes initialized\n"
     ]
    }
   ],
   "source": [
    "points_vars = []\n",
    "width_vars = []\n",
    "color_vars = []\n",
    "\n",
    "strokes = []\n",
    "\n",
    "for i in range(num_paths):\n",
    "    xyz = []\n",
    "    rgba = []\n",
    "    ws = []\n",
    "    \n",
    "    p0 = 2*np.random.rand(3)-1\n",
    "    w0 = np.random.rand()*0.1+0.02\n",
    "    rgb = np.random.rand(4)\n",
    "    \n",
    "    for j in range(4):\n",
    "        xyz.append(p0)\n",
    "        ws.append(w0)\n",
    "        rgba.append(rgb)\n",
    "        \n",
    "        p0 = p0 + 0.3 * (2*np.random.rand(3)-1)\n",
    "    \n",
    "    xyz = torch.FloatTensor(np.array(xyz)).to(device)\n",
    "    xyz.requires_grad = True\n",
    "    points_vars.append(xyz)\n",
    "    \n",
    "    rgba = torch.FloatTensor(np.array(rgba)).to(device)\n",
    "    rgba.requires_grad = True\n",
    "    color_vars.append(rgba)\n",
    "\n",
    "    ws = torch.FloatTensor(np.array(ws)).to(device)\n",
    "    ws.requires_grad = True\n",
    "    width_vars.append(ws)\n",
    "    \n",
    "    strokes.append([xyz, rgba, ws])\n",
    "\n",
    "print(\"Strokes initialized\")\n",
    "# Optimize\n",
    "points_optim = torch.optim.Adam(points_vars, lr=0.05)\n",
    "width_optim = torch.optim.Adam(width_vars, lr=0.02)\n",
    "color_optim = torch.optim.Adam(color_vars, lr=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train strokes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "render loss: -0.24462890625\n",
      "iteration: 0\n",
      "render loss: -0.2548828125\n",
      "iteration: 1\n",
      "render loss: -0.271240234375\n",
      "iteration: 2\n",
      "render loss: -0.27978515625\n",
      "iteration: 3\n",
      "render loss: -0.28564453125\n",
      "iteration: 4\n",
      "render loss: -0.29052734375\n",
      "iteration: 5\n",
      "render loss: -0.29296875\n",
      "iteration: 6\n",
      "render loss: -0.298828125\n",
      "iteration: 7\n",
      "render loss: -0.288330078125\n",
      "iteration: 8\n",
      "render loss: -0.30126953125\n",
      "iteration: 9\n",
      "render loss: -0.296630859375\n",
      "iteration: 10\n",
      "render loss: -0.29833984375\n",
      "iteration: 11\n",
      "render loss: -0.29345703125\n",
      "iteration: 12\n",
      "render loss: -0.305419921875\n",
      "iteration: 13\n",
      "render loss: -0.310302734375\n",
      "iteration: 14\n",
      "render loss: -0.31103515625\n",
      "iteration: 15\n",
      "render loss: -0.31787109375\n",
      "iteration: 16\n",
      "render loss: -0.29931640625\n",
      "iteration: 17\n",
      "render loss: -0.3056640625\n",
      "iteration: 18\n",
      "render loss: -0.31689453125\n",
      "iteration: 19\n",
      "render loss: -0.326171875\n",
      "iteration: 20\n",
      "render loss: -0.33056640625\n",
      "iteration: 21\n",
      "render loss: -0.322509765625\n",
      "iteration: 22\n",
      "render loss: -0.328369140625\n",
      "iteration: 23\n",
      "render loss: -0.338134765625\n",
      "iteration: 24\n",
      "render loss: -0.3349609375\n",
      "iteration: 25\n",
      "render loss: -0.33642578125\n",
      "iteration: 26\n",
      "render loss: -0.332763671875\n",
      "iteration: 27\n",
      "render loss: -0.33642578125\n",
      "iteration: 28\n",
      "render loss: -0.3427734375\n",
      "iteration: 29\n",
      "render loss: -0.34033203125\n",
      "iteration: 30\n",
      "render loss: -0.33544921875\n",
      "iteration: 31\n",
      "render loss: -0.3369140625\n",
      "iteration: 32\n",
      "render loss: -0.333984375\n",
      "iteration: 33\n",
      "render loss: -0.338134765625\n",
      "iteration: 34\n",
      "render loss: -0.335205078125\n",
      "iteration: 35\n",
      "render loss: -0.338134765625\n",
      "iteration: 36\n",
      "render loss: -0.34326171875\n",
      "iteration: 37\n",
      "render loss: -0.34619140625\n",
      "iteration: 38\n",
      "render loss: -0.331787109375\n",
      "iteration: 39\n",
      "render loss: -0.33251953125\n",
      "iteration: 40\n",
      "render loss: -0.3388671875\n",
      "iteration: 41\n",
      "render loss: -0.33251953125\n",
      "iteration: 42\n",
      "render loss: -0.332763671875\n",
      "iteration: 43\n",
      "render loss: -0.338623046875\n",
      "iteration: 44\n",
      "render loss: -0.33837890625\n",
      "iteration: 45\n",
      "render loss: -0.322021484375\n",
      "iteration: 46\n",
      "render loss: -0.326171875\n",
      "iteration: 47\n",
      "render loss: -0.32421875\n",
      "iteration: 48\n",
      "render loss: -0.345947265625\n",
      "iteration: 49\n",
      "render loss: -0.337646484375\n",
      "iteration: 50\n",
      "render loss: -0.3369140625\n",
      "iteration: 51\n",
      "render loss: -0.33544921875\n",
      "iteration: 52\n",
      "render loss: -0.341796875\n",
      "iteration: 53\n",
      "render loss: -0.3330078125\n",
      "iteration: 54\n",
      "render loss: -0.32177734375\n",
      "iteration: 55\n",
      "render loss: -0.333740234375\n",
      "iteration: 56\n",
      "render loss: -0.33251953125\n",
      "iteration: 57\n",
      "render loss: -0.33349609375\n",
      "iteration: 58\n",
      "render loss: -0.336181640625\n",
      "iteration: 59\n",
      "render loss: -0.33203125\n",
      "iteration: 60\n",
      "render loss: -0.341064453125\n",
      "iteration: 61\n",
      "render loss: -0.344482421875\n",
      "iteration: 62\n",
      "render loss: -0.34130859375\n",
      "iteration: 63\n",
      "render loss: -0.33740234375\n",
      "iteration: 64\n",
      "render loss: -0.345458984375\n",
      "iteration: 65\n",
      "render loss: -0.34228515625\n",
      "iteration: 66\n",
      "render loss: -0.337158203125\n",
      "iteration: 67\n",
      "render loss: -0.343994140625\n",
      "iteration: 68\n",
      "render loss: -0.34033203125\n",
      "iteration: 69\n",
      "render loss: -0.345703125\n",
      "iteration: 70\n",
      "render loss: -0.34716796875\n",
      "iteration: 71\n",
      "render loss: -0.337158203125\n",
      "iteration: 72\n",
      "render loss: -0.323974609375\n",
      "iteration: 73\n",
      "render loss: -0.33740234375\n",
      "iteration: 74\n",
      "render loss: -0.3466796875\n",
      "iteration: 75\n",
      "render loss: -0.34130859375\n",
      "iteration: 76\n",
      "render loss: -0.313232421875\n",
      "iteration: 77\n",
      "render loss: -0.33349609375\n",
      "iteration: 78\n",
      "render loss: -0.3349609375\n",
      "iteration: 79\n",
      "render loss: -0.335693359375\n",
      "iteration: 80\n",
      "render loss: -0.345458984375\n",
      "iteration: 81\n",
      "render loss: -0.333740234375\n",
      "iteration: 82\n",
      "render loss: -0.329345703125\n",
      "iteration: 83\n",
      "render loss: -0.32275390625\n",
      "iteration: 84\n",
      "render loss: -0.3330078125\n",
      "iteration: 85\n",
      "render loss: -0.34375\n",
      "iteration: 86\n",
      "render loss: -0.332275390625\n",
      "iteration: 87\n",
      "render loss: -0.344970703125\n",
      "iteration: 88\n",
      "render loss: -0.33056640625\n",
      "iteration: 89\n",
      "render loss: -0.345458984375\n",
      "iteration: 90\n",
      "render loss: -0.34423828125\n",
      "iteration: 91\n",
      "render loss: -0.34912109375\n",
      "iteration: 92\n",
      "render loss: -0.33544921875\n",
      "iteration: 93\n",
      "render loss: -0.340087890625\n",
      "iteration: 94\n",
      "render loss: -0.337890625\n",
      "iteration: 95\n",
      "render loss: -0.328125\n",
      "iteration: 96\n",
      "render loss: -0.330078125\n",
      "iteration: 97\n",
      "render loss: -0.330078125\n",
      "iteration: 98\n",
      "render loss: -0.339111328125\n",
      "iteration: 99\n",
      "render loss: -0.334228515625\n",
      "iteration: 100\n",
      "render loss: -0.340576171875\n",
      "iteration: 101\n",
      "render loss: -0.34228515625\n",
      "iteration: 102\n",
      "render loss: -0.3359375\n",
      "iteration: 103\n",
      "render loss: -0.33154296875\n",
      "iteration: 104\n",
      "render loss: -0.331787109375\n",
      "iteration: 105\n",
      "render loss: -0.33154296875\n",
      "iteration: 106\n",
      "render loss: -0.33642578125\n",
      "iteration: 107\n",
      "render loss: -0.34033203125\n",
      "iteration: 108\n",
      "render loss: -0.342529296875\n",
      "iteration: 109\n",
      "render loss: -0.34228515625\n",
      "iteration: 110\n",
      "render loss: -0.331298828125\n",
      "iteration: 111\n",
      "render loss: -0.325439453125\n",
      "iteration: 112\n",
      "render loss: -0.32568359375\n",
      "iteration: 113\n",
      "render loss: -0.326904296875\n",
      "iteration: 114\n",
      "render loss: -0.338623046875\n",
      "iteration: 115\n",
      "render loss: -0.339111328125\n",
      "iteration: 116\n",
      "render loss: -0.335693359375\n",
      "iteration: 117\n",
      "render loss: -0.337158203125\n",
      "iteration: 118\n",
      "render loss: -0.33544921875\n",
      "iteration: 119\n",
      "render loss: -0.34619140625\n",
      "iteration: 120\n",
      "render loss: -0.345947265625\n",
      "iteration: 121\n",
      "render loss: -0.339599609375\n",
      "iteration: 122\n",
      "render loss: -0.341064453125\n",
      "iteration: 123\n",
      "render loss: -0.340087890625\n",
      "iteration: 124\n",
      "render loss: -0.343994140625\n",
      "iteration: 125\n",
      "render loss: -0.351318359375\n",
      "iteration: 126\n",
      "render loss: -0.34423828125\n",
      "iteration: 127\n",
      "render loss: -0.341064453125\n",
      "iteration: 128\n",
      "render loss: -0.348876953125\n",
      "iteration: 129\n",
      "render loss: -0.343994140625\n",
      "iteration: 130\n",
      "render loss: -0.34326171875\n",
      "iteration: 131\n",
      "render loss: -0.345947265625\n",
      "iteration: 132\n",
      "render loss: -0.351318359375\n",
      "iteration: 133\n",
      "render loss: -0.34912109375\n",
      "iteration: 134\n",
      "render loss: -0.35009765625\n",
      "iteration: 135\n",
      "render loss: -0.33984375\n",
      "iteration: 136\n",
      "render loss: -0.348876953125\n",
      "iteration: 137\n",
      "render loss: -0.343017578125\n",
      "iteration: 138\n",
      "render loss: -0.345458984375\n",
      "iteration: 139\n",
      "render loss: -0.32568359375\n",
      "iteration: 140\n",
      "render loss: -0.353515625\n",
      "iteration: 141\n",
      "render loss: -0.3427734375\n",
      "iteration: 142\n",
      "render loss: -0.3486328125\n",
      "iteration: 143\n",
      "render loss: -0.351318359375\n",
      "iteration: 144\n",
      "render loss: -0.35009765625\n",
      "iteration: 145\n",
      "render loss: -0.353515625\n",
      "iteration: 146\n",
      "render loss: -0.352294921875\n",
      "iteration: 147\n",
      "render loss: -0.3466796875\n",
      "iteration: 148\n",
      "render loss: -0.347412109375\n",
      "iteration: 149\n",
      "render loss: -0.349853515625\n",
      "iteration: 150\n",
      "render loss: -0.34375\n",
      "iteration: 151\n",
      "render loss: -0.327880859375\n",
      "iteration: 152\n",
      "render loss: -0.331787109375\n",
      "iteration: 153\n",
      "render loss: -0.3486328125\n",
      "iteration: 154\n",
      "render loss: -0.3310546875\n",
      "iteration: 155\n",
      "render loss: -0.31640625\n",
      "iteration: 156\n",
      "render loss: -0.331298828125\n",
      "iteration: 157\n",
      "render loss: -0.340576171875\n",
      "iteration: 158\n",
      "render loss: -0.337890625\n",
      "iteration: 159\n",
      "render loss: -0.347412109375\n",
      "iteration: 160\n",
      "render loss: -0.349853515625\n",
      "iteration: 161\n",
      "render loss: -0.346435546875\n",
      "iteration: 162\n",
      "render loss: -0.34912109375\n",
      "iteration: 163\n",
      "render loss: -0.333984375\n",
      "iteration: 164\n",
      "render loss: -0.3349609375\n",
      "iteration: 165\n",
      "render loss: -0.330078125\n",
      "iteration: 166\n",
      "render loss: -0.32763671875\n",
      "iteration: 167\n",
      "render loss: -0.34033203125\n",
      "iteration: 168\n",
      "render loss: -0.3359375\n",
      "iteration: 169\n",
      "render loss: -0.34228515625\n",
      "iteration: 170\n",
      "render loss: -0.330078125\n",
      "iteration: 171\n",
      "render loss: -0.339599609375\n",
      "iteration: 172\n",
      "render loss: -0.348876953125\n",
      "iteration: 173\n",
      "render loss: -0.34521484375\n",
      "iteration: 174\n",
      "render loss: -0.336181640625\n",
      "iteration: 175\n",
      "render loss: -0.345703125\n",
      "iteration: 176\n",
      "render loss: -0.353271484375\n",
      "iteration: 177\n",
      "render loss: -0.33837890625\n",
      "iteration: 178\n",
      "render loss: -0.3330078125\n",
      "iteration: 179\n",
      "render loss: -0.3486328125\n",
      "iteration: 180\n",
      "render loss: -0.341796875\n",
      "iteration: 181\n",
      "render loss: -0.34375\n",
      "iteration: 182\n",
      "render loss: -0.34912109375\n",
      "iteration: 183\n",
      "render loss: -0.353515625\n",
      "iteration: 184\n",
      "render loss: -0.3427734375\n",
      "iteration: 185\n",
      "render loss: -0.353759765625\n",
      "iteration: 186\n",
      "render loss: -0.343994140625\n",
      "iteration: 187\n",
      "render loss: -0.335693359375\n",
      "iteration: 188\n",
      "render loss: -0.342529296875\n",
      "iteration: 189\n",
      "render loss: -0.34033203125\n",
      "iteration: 190\n",
      "render loss: -0.3369140625\n",
      "iteration: 191\n",
      "render loss: -0.3427734375\n",
      "iteration: 192\n",
      "render loss: -0.3505859375\n",
      "iteration: 193\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "render loss: -0.3271484375\n",
      "iteration: 194\n",
      "render loss: -0.338623046875\n",
      "iteration: 195\n",
      "render loss: -0.34326171875\n",
      "iteration: 196\n",
      "render loss: -0.33642578125\n",
      "iteration: 197\n",
      "render loss: -0.34033203125\n",
      "iteration: 198\n",
      "render loss: -0.33984375\n",
      "iteration: 199\n",
      "render loss: -0.340576171875\n",
      "iteration: 200\n",
      "render loss: -0.343994140625\n",
      "iteration: 201\n",
      "render loss: -0.3349609375\n",
      "iteration: 202\n",
      "render loss: -0.34716796875\n",
      "iteration: 203\n",
      "render loss: -0.34130859375\n",
      "iteration: 204\n",
      "render loss: -0.342041015625\n",
      "iteration: 205\n",
      "render loss: -0.34765625\n",
      "iteration: 206\n",
      "render loss: -0.34814453125\n",
      "iteration: 207\n",
      "render loss: -0.348388671875\n",
      "iteration: 208\n",
      "render loss: -0.344482421875\n",
      "iteration: 209\n",
      "render loss: -0.347900390625\n",
      "iteration: 210\n",
      "render loss: -0.346923828125\n",
      "iteration: 211\n",
      "render loss: -0.352783203125\n",
      "iteration: 212\n",
      "render loss: -0.3486328125\n",
      "iteration: 213\n",
      "render loss: -0.350341796875\n",
      "iteration: 214\n",
      "render loss: -0.35693359375\n",
      "iteration: 215\n",
      "render loss: -0.353759765625\n",
      "iteration: 216\n",
      "render loss: -0.347900390625\n",
      "iteration: 217\n",
      "render loss: -0.356689453125\n",
      "iteration: 218\n",
      "render loss: -0.349609375\n",
      "iteration: 219\n",
      "render loss: -0.3505859375\n",
      "iteration: 220\n",
      "render loss: -0.353759765625\n",
      "iteration: 221\n",
      "render loss: -0.3525390625\n",
      "iteration: 222\n",
      "render loss: -0.35498046875\n",
      "iteration: 223\n",
      "render loss: -0.34716796875\n",
      "iteration: 224\n",
      "render loss: -0.3525390625\n",
      "iteration: 225\n",
      "render loss: -0.357666015625\n",
      "iteration: 226\n",
      "render loss: -0.3544921875\n",
      "iteration: 227\n",
      "render loss: -0.35888671875\n",
      "iteration: 228\n",
      "render loss: -0.359619140625\n",
      "iteration: 229\n",
      "render loss: -0.360595703125\n",
      "iteration: 230\n",
      "render loss: -0.3544921875\n",
      "iteration: 231\n",
      "render loss: -0.35986328125\n",
      "iteration: 232\n",
      "render loss: -0.363525390625\n",
      "iteration: 233\n",
      "render loss: -0.35498046875\n",
      "iteration: 234\n",
      "render loss: -0.363037109375\n",
      "iteration: 235\n",
      "render loss: -0.361328125\n",
      "iteration: 236\n",
      "render loss: -0.3564453125\n",
      "iteration: 237\n",
      "render loss: -0.359619140625\n",
      "iteration: 238\n",
      "render loss: -0.35986328125\n",
      "iteration: 239\n",
      "render loss: -0.354248046875\n",
      "iteration: 240\n",
      "render loss: -0.36669921875\n",
      "iteration: 241\n",
      "render loss: -0.3681640625\n",
      "iteration: 242\n",
      "render loss: -0.363037109375\n",
      "iteration: 243\n",
      "render loss: -0.365966796875\n",
      "iteration: 244\n",
      "render loss: -0.364501953125\n",
      "iteration: 245\n",
      "render loss: -0.365234375\n",
      "iteration: 246\n",
      "render loss: -0.36572265625\n",
      "iteration: 247\n",
      "render loss: -0.363525390625\n",
      "iteration: 248\n",
      "render loss: -0.3642578125\n",
      "iteration: 249\n",
      "render loss: -0.360595703125\n",
      "iteration: 250\n",
      "render loss: -0.361328125\n",
      "iteration: 251\n",
      "render loss: -0.34765625\n",
      "iteration: 252\n",
      "render loss: -0.35791015625\n",
      "iteration: 253\n",
      "render loss: -0.3505859375\n",
      "iteration: 254\n",
      "render loss: -0.359619140625\n",
      "iteration: 255\n",
      "render loss: -0.35888671875\n",
      "iteration: 256\n",
      "render loss: -0.364990234375\n",
      "iteration: 257\n",
      "render loss: -0.365966796875\n",
      "iteration: 258\n",
      "render loss: -0.35693359375\n",
      "iteration: 259\n",
      "render loss: -0.36181640625\n",
      "iteration: 260\n",
      "render loss: -0.36279296875\n",
      "iteration: 261\n",
      "render loss: -0.361572265625\n",
      "iteration: 262\n",
      "render loss: -0.359619140625\n",
      "iteration: 263\n",
      "render loss: -0.365478515625\n",
      "iteration: 264\n",
      "render loss: -0.35693359375\n",
      "iteration: 265\n",
      "render loss: -0.37255859375\n",
      "iteration: 266\n",
      "render loss: -0.360107421875\n",
      "iteration: 267\n",
      "render loss: -0.375\n",
      "iteration: 268\n",
      "render loss: -0.367919921875\n",
      "iteration: 269\n",
      "render loss: -0.36669921875\n",
      "iteration: 270\n",
      "render loss: -0.37255859375\n",
      "iteration: 271\n",
      "render loss: -0.372314453125\n",
      "iteration: 272\n",
      "render loss: -0.367431640625\n",
      "iteration: 273\n",
      "render loss: -0.373779296875\n",
      "iteration: 274\n",
      "render loss: -0.35546875\n",
      "iteration: 275\n",
      "render loss: -0.362060546875\n",
      "iteration: 276\n",
      "render loss: -0.362060546875\n",
      "iteration: 277\n",
      "render loss: -0.37060546875\n",
      "iteration: 278\n",
      "render loss: -0.36083984375\n",
      "iteration: 279\n",
      "render loss: -0.364990234375\n",
      "iteration: 280\n",
      "render loss: -0.373046875\n",
      "iteration: 281\n",
      "render loss: -0.35302734375\n",
      "iteration: 282\n",
      "render loss: -0.36279296875\n",
      "iteration: 283\n",
      "render loss: -0.36767578125\n",
      "iteration: 284\n",
      "render loss: -0.37158203125\n",
      "iteration: 285\n",
      "render loss: -0.368408203125\n",
      "iteration: 286\n",
      "render loss: -0.368896484375\n",
      "iteration: 287\n",
      "render loss: -0.368408203125\n",
      "iteration: 288\n",
      "render loss: -0.370361328125\n",
      "iteration: 289\n",
      "render loss: -0.37158203125\n",
      "iteration: 290\n",
      "render loss: -0.370849609375\n",
      "iteration: 291\n",
      "render loss: -0.37158203125\n",
      "iteration: 292\n",
      "render loss: -0.37060546875\n",
      "iteration: 293\n",
      "render loss: -0.377685546875\n",
      "iteration: 294\n",
      "render loss: -0.378173828125\n",
      "iteration: 295\n",
      "render loss: -0.372802734375\n",
      "iteration: 296\n",
      "render loss: -0.3740234375\n",
      "iteration: 297\n",
      "render loss: -0.37890625\n",
      "iteration: 298\n",
      "render loss: -0.373291015625\n",
      "iteration: 299\n",
      "render loss: -0.3759765625\n",
      "iteration: 300\n",
      "render loss: -0.3681640625\n",
      "iteration: 301\n",
      "render loss: -0.367431640625\n",
      "iteration: 302\n",
      "render loss: -0.373046875\n",
      "iteration: 303\n",
      "render loss: -0.376220703125\n",
      "iteration: 304\n",
      "render loss: -0.370361328125\n",
      "iteration: 305\n",
      "render loss: -0.375244140625\n",
      "iteration: 306\n",
      "render loss: -0.372314453125\n",
      "iteration: 307\n",
      "render loss: -0.3759765625\n",
      "iteration: 308\n",
      "render loss: -0.375\n",
      "iteration: 309\n",
      "render loss: -0.37744140625\n",
      "iteration: 310\n",
      "render loss: -0.3740234375\n",
      "iteration: 311\n",
      "render loss: -0.3798828125\n",
      "iteration: 312\n",
      "render loss: -0.37939453125\n",
      "iteration: 313\n",
      "render loss: -0.375732421875\n",
      "iteration: 314\n",
      "render loss: -0.3720703125\n",
      "iteration: 315\n",
      "render loss: -0.3798828125\n",
      "iteration: 316\n",
      "render loss: -0.381591796875\n",
      "iteration: 317\n",
      "render loss: -0.376953125\n",
      "iteration: 318\n",
      "render loss: -0.376220703125\n",
      "iteration: 319\n",
      "render loss: -0.377197265625\n",
      "iteration: 320\n",
      "render loss: -0.37158203125\n",
      "iteration: 321\n",
      "render loss: -0.36865234375\n",
      "iteration: 322\n",
      "render loss: -0.37109375\n",
      "iteration: 323\n",
      "render loss: -0.373291015625\n",
      "iteration: 324\n",
      "render loss: -0.37646484375\n",
      "iteration: 325\n",
      "render loss: -0.379150390625\n",
      "iteration: 326\n",
      "render loss: -0.380126953125\n",
      "iteration: 327\n",
      "render loss: -0.3759765625\n",
      "iteration: 328\n",
      "render loss: -0.374267578125\n",
      "iteration: 329\n",
      "render loss: -0.3740234375\n",
      "iteration: 330\n",
      "render loss: -0.37451171875\n",
      "iteration: 331\n",
      "render loss: -0.376953125\n",
      "iteration: 332\n",
      "render loss: -0.373046875\n",
      "iteration: 333\n",
      "render loss: -0.37841796875\n",
      "iteration: 334\n",
      "render loss: -0.376953125\n",
      "iteration: 335\n",
      "render loss: -0.37744140625\n",
      "iteration: 336\n",
      "render loss: -0.374267578125\n",
      "iteration: 337\n",
      "render loss: -0.370849609375\n",
      "iteration: 338\n",
      "render loss: -0.3720703125\n",
      "iteration: 339\n",
      "render loss: -0.375\n",
      "iteration: 340\n",
      "render loss: -0.370361328125\n",
      "iteration: 341\n",
      "render loss: -0.374267578125\n",
      "iteration: 342\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-7-a073e52538ce>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     54\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     55\u001b[0m         \u001b[0;31m# Backpropagate the gradients.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 56\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     57\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     58\u001b[0m     \u001b[0;31m# Take a gradient descent step.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    243\u001b[0m                 \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    244\u001b[0m                 inputs=inputs)\n\u001b[0;32m--> 245\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    246\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    247\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    143\u001b[0m         \u001b[0mretain_graph\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    144\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 145\u001b[0;31m     Variable._execution_engine.run_backward(\n\u001b[0m\u001b[1;32m    146\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    147\u001b[0m         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# Adam iterations.\n",
    "for t in range(args.num_iter):\n",
    "    points_optim.zero_grad()\n",
    "    width_optim.zero_grad()\n",
    "    color_optim.zero_grad()\n",
    "    \n",
    "    if t==200: # Some learning rate decay to help a drawing settle\n",
    "        for p in points_optim.param_groups:\n",
    "            p['lr'] *= 0.2\n",
    "\n",
    "        for p in width_optim.param_groups:\n",
    "            p['lr'] *= 0.2\n",
    "\n",
    "        for p in color_optim.param_groups:\n",
    "            p['lr'] *= 0.2    \n",
    "    \n",
    "    for n in range(4):\n",
    "        if n==3: # Use a fixed reference for at least one, for visualization\n",
    "            cam_pos = np.array([1, 0.1, 0])\n",
    "        else:\n",
    "            cam_pos = np.array([np.random.randn(), 0.2*np.random.rand(), np.random.randn()])\n",
    "        \n",
    "        theta = atan2(cam_pos[2], cam_pos[0])\n",
    "        \n",
    "        theta = (theta+2*3.14159) % (2*3.14159)\n",
    "        \n",
    "        if theta>0 and theta<=2*3.14159/3:\n",
    "            features = text_features1\n",
    "        elif theta<=4*3.14159/3:\n",
    "            features = text_features2\n",
    "        else:\n",
    "            features = text_features3\n",
    "            \n",
    "        cam_pos = F.normalize(torch.FloatTensor(cam_pos).to(device), dim=0)*3\n",
    "        cam_targ = torch.zeros(3).to(device)\n",
    "\n",
    "        img = render(cam_pos, cam_targ, strokes)\n",
    "        \n",
    "        # Save the intermediate render.\n",
    "        if n==3:\n",
    "            im = Image.fromarray(np.clip(255*img.cpu().detach().numpy(),0,255).astype(np.uint8))\n",
    "            im.save(\"tmp/inter_%.6d.png\" % t)\n",
    "        \n",
    "        # Convert img from HWC to NCHW\n",
    "        img = img.unsqueeze(0)\n",
    "        img = img.permute(0, 3, 1, 2) # NHWC -> NCHW\n",
    "\n",
    "        loss = 0\n",
    "        for m in range(8):\n",
    "            img_aug = augment_trans(img)\n",
    "            image_features = model.encode_image(img_aug)\n",
    "            loss -= torch.cosine_similarity(features, image_features, dim=1)/8.0\n",
    "            \n",
    "\n",
    "        # Backpropagate the gradients.\n",
    "        loss.backward()\n",
    "\n",
    "    # Take a gradient descent step.\n",
    "    points_optim.step()\n",
    "    width_optim.step()\n",
    "    color_optim.step()\n",
    "    for s in strokes:\n",
    "        s[2].data.clamp_(0.02, max_width)\n",
    "        s[1].data.clamp_(0.0, 1.0)\n",
    "        s[0].data.clamp_(-1.0, 1.0)\n",
    "    \n",
    "    if t % 1 == 0:\n",
    "        # show_img(torch.cat([img.detach(), img_aug.detach()], axis=3).cpu().numpy()[0])\n",
    "        print('render loss:', loss.item())\n",
    "        print('iteration:', t)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Render spinning view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import *\n",
    "\n",
    "frame = 0\n",
    "\n",
    "for theta in np.arange(0,2*3.1415,0.05):\n",
    "    with torch.no_grad():\n",
    "        cam_pos = torch.cuda.FloatTensor(np.array([3*cos(theta), 0.2, 3*sin(theta)]))\n",
    "        cam_targ = torch.zeros(3).to(device)\n",
    "    \n",
    "        img = render(cam_pos, cam_targ, strokes)\n",
    "\n",
    "        # Save the intermediate render.\n",
    "        im = Image.fromarray(np.clip(255*img.cpu().detach().numpy(),0,255).astype(np.uint8))\n",
    "        im.save(\"output/%.6d.png\" % frame)\n",
    "        \n",
    "        frame += 1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save strokes for other applications/re-rendering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_filename = \"a_bunch_of_strokes.pth\"\n",
    "\n",
    "torch.save(strokes, open(output_filename,\"wb\"))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "CLIPDraw AI Telephone (Cleaned).ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}