Revert "remove restriction to run conv_norm in fp32"

This reverts commit cec5928.

Author: NouamaneTazi

src/diffusers/models/resnet.py

@@ -332,8 +332,7 @@ def forward(self, x, temb):
 
         # make sure hidden states is in float32
         # when running in half-precision
-        hidden_states = self.norm1(hidden_states).type(hidden_states.dtype)
-        # hidden_states = self.norm1.float()(hidden_states.float()).type(hidden_states.dtype)
+        hidden_states = self.norm1.float()(hidden_states.float()).type(hidden_states.dtype)
         hidden_states = self.nonlinearity(hidden_states)
 
         if self.upsample is not None:
@@ -351,8 +350,7 @@ def forward(self, x, temb):
 
         # make sure hidden states is in float32
         # when running in half-precision
-        hidden_states = self.norm2(hidden_states).type(hidden_states.dtype)
-        # hidden_states = self.norm2.float()(hidden_states.float()).type(hidden_states.dtype)
+        hidden_states = self.norm2.float()(hidden_states.float()).type(hidden_states.dtype)
         hidden_states = self.nonlinearity(hidden_states)
 
         hidden_states = self.dropout(hidden_states)

src/diffusers/models/unet_2d_condition.py

@@ -261,8 +261,7 @@ def forward(
         # 6. post-process
         # make sure hidden states is in float32
         # when running in half-precision
-        # sample = self.conv_norm_out.float()(sample.float()).type(sample.dtype)
-        sample = self.conv_norm_out(sample).type(sample.dtype)
+        sample = self.conv_norm_out.float()(sample.float()).type(sample.dtype)
         sample = self.conv_act(sample)
         sample = self.conv_out(sample)