fix dynamicrafter notebook (#2376)

openvinotoolkit · Sep 10, 2024 · 406f894 · 406f894
1 parent e57fb8a
commit 406f894
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/notebooks/dynamicrafter-animating-images/dynamicrafter-animating-images.ipynb b/notebooks/dynamicrafter-animating-images/dynamicrafter-animating-images.ipynb
@@ -559,8 +559,9 @@
  "import open_clip\n",
  "\n",
  "\n",
- "class CondStageModelWrapper:\n",
+ "class CondStageModelWrapper(torch.nn.Module):\n",
  " def __init__(self, cond_stage_model):\n",
+ " super().__init__()\n",
  " self.cond_stage_model = cond_stage_model\n",
  "\n",
  " def encode(self, tokens):\n",
@@ -571,8 +572,9 @@
  " return torch.from_numpy(outs)\n",
  "\n",
  "\n",
- "class EncoderFirstStageModelWrapper:\n",
+ "class EncoderFirstStageModelWrapper(torch.nn.Module):\n",
  " def __init__(self, encode_first_stage):\n",
+ " super().__init__()\n",
  " self.encode_first_stage = encode_first_stage\n",
  "\n",
  " def forward(self, x):\n",
@@ -584,8 +586,9 @@
  " return self.forward(*args, **kwargs)\n",
  "\n",
  "\n",
- "class EmbedderWrapper:\n",
+ "class EmbedderWrapper(torch.nn.Module):\n",
  " def __init__(self, embedder):\n",
+ " super().__init__()\n",
  " self.embedder = embedder\n",
  "\n",
  " def forward(self, x):\n",
@@ -597,8 +600,9 @@
  " return self.forward(*args, **kwargs)\n",
  "\n",
  "\n",
- "class CModelWrapper:\n",
+ "class CModelWrapper(torch.nn.Module):\n",
  " def __init__(self, diffusion_model, out_channels):\n",
+ " super().__init__()\n",
  " self.diffusion_model = diffusion_model\n",
  " self.out_channels = out_channels\n",
  "\n",
@@ -617,8 +621,9 @@
  " return self.forward(*args, **kwargs)\n",
  "\n",
  "\n",
- "class DecoderFirstStageModelWrapper:\n",
+ "class DecoderFirstStageModelWrapper(torch.nn.Module):\n",
  " def __init__(self, decoder_first_stage):\n",
+ " super().__init__()\n",
  " self.decoder_first_stage = decoder_first_stage\n",
  "\n",
  " def forward(self, x):\n",
@@ -1875,7 +1880,7 @@
  "name": "python",
  "nbconvert_exporter": "python",
  "pygments_lexer": "ipython3",
- "version": "3.8.10"
+ "version": "3.11.4"
  },
  "openvino_notebooks": {
  "imageUrl": "https://github.com/Doubiiu/DynamiCrafter/blob/main/assets/showcase/guitar0.gif?raw=true",