From 8bb362a853d38bd0eb3e6e786561af7d70bbb938 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Wed, 14 Aug 2024 11:49:41 +0800
Subject: [PATCH] Update openaimodel_diffree.py

---
 .../ldmx/modules/diffusionmodules/openaimodel_diffree.py       | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/stable_diffusion/ldmx/modules/diffusionmodules/openaimodel_diffree.py b/stable_diffusion/ldmx/modules/diffusionmodules/openaimodel_diffree.py
index af2aaf6..6d2f093 100644
--- a/stable_diffusion/ldmx/modules/diffusionmodules/openaimodel_diffree.py
+++ b/stable_diffusion/ldmx/modules/diffusionmodules/openaimodel_diffree.py
@@ -818,7 +818,6 @@ class UNetModel(nn.Module):
 
 class OMPModule(nn.Module):
     """
-    The full UNet model with attention and timestep embedding.
     :param in_channels: channels in the input Tensor.
     :param model_channels: base channel count for the model.
     :param out_channels: channels in the output Tensor.
@@ -828,7 +827,7 @@ class OMPModule(nn.Module):
         For example, if this contains 4, then at 4x downsampling, attention
         will be used.
     :param dropout: the dropout probability.
-    :param channel_mult: channel multiplier for each level of the UNet.
+    :param channel_mult: channel multiplier for each level of the OMPModule.
     :param conv_resample: if True, use learned convolutions for upsampling and
         downsampling.
     :param dims: determines if the signal is 1D, 2D, or 3D.