Merge pull request #1995 from dreavjr:deepdream-correct-shape

PiperOrigin-RevId: 421487418

Author: copybara-github

site/en/tutorials/generative/deepdream.ipynb

@@ -512,19 +512,20 @@
         "  @tf.function(\n",
         "      input_signature=(\n",
         "        tf.TensorSpec(shape=[None,None,3], dtype=tf.float32),\n",
+        "        tf.TensorSpec(shape=[2], dtype=tf.int32),\n",
         "        tf.TensorSpec(shape=[], dtype=tf.int32),)\n",
         "  )\n",
-        "  def __call__(self, img, tile_size=512):\n",
+        "  def __call__(self, img, img_size, tile_size=512):\n",
         "    shift, img_rolled = random_roll(img, tile_size)\n",
         "\n",
         "    # Initialize the image gradients to zero.\n",
         "    gradients = tf.zeros_like(img_rolled)\n",
         "    \n",
         "    # Skip the last tile, unless there's only one tile.\n",
-        "    xs = tf.range(0, img_rolled.shape[0], tile_size)[:-1]\n",
+        "    xs = tf.range(0, img_size[1], tile_size)[:-1]\n",
         "    if not tf.cast(len(xs), bool):\n",
         "      xs = tf.constant([0])\n",
-        "    ys = tf.range(0, img_rolled.shape[1], tile_size)[:-1]\n",
+        "    ys = tf.range(0, img_size[0], tile_size)[:-1]\n",
         "    if not tf.cast(len(ys), bool):\n",
         "      ys = tf.constant([0])\n",
         "\n",
@@ -537,7 +538,7 @@
         "          tape.watch(img_rolled)\n",
         "\n",
         "          # Extract a tile out of the image.\n",
-        "          img_tile = img_rolled[x:x+tile_size, y:y+tile_size]\n",
+        "          img_tile = img_rolled[y:y+tile_size, x:x+tile_size]\n",
         "          loss = calc_loss(img_tile, self.model)\n",
         "\n",
         "        # Update the image gradients for this tile.\n",
@@ -591,10 +592,11 @@
         "  for octave in octaves:\n",
         "    # Scale the image based on the octave\n",
         "    new_size = tf.cast(tf.convert_to_tensor(base_shape[:-1]), tf.float32)*(octave_scale**octave)\n",
-        "    img = tf.image.resize(img, tf.cast(new_size, tf.int32))\n",
+        "    new_size = tf.cast(new_size, tf.int32)\n",
+        "    img = tf.image.resize(img, new_size)\n",
         "\n",
         "    for step in range(steps_per_octave):\n",
-        "      gradients = get_tiled_gradients(img)\n",
+        "      gradients = get_tiled_gradients(img, new_size)\n",
         "      img = img + gradients*step_size\n",
         "      img = tf.clip_by_value(img, -1, 1)\n",
         "\n",