RL Cleanup v2

scripts/convert_models_diffuser_to_diffusers.py

@@ -49,6 +49,7 @@ def value_function():
         down_block_types=("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"),
         up_block_types=(),
         out_block_type="ValueFunction",
+        mid_block_type="ValueFunctionMidBlock1D",
         block_out_channels=(32, 64, 128, 256),
         layers_per_block=1,
         always_downsample=True,

src/diffusers/models/unet_1d_blocks.py

@@ -173,6 +173,26 @@ class UpBlock1DNoSkip(nn.Module):
     pass
 
 
+class ValueFunctionMidBlock1D(nn.Module):
+    def __init__(self, in_channels, out_channels, embed_dim):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.embed_dim = embed_dim
+
+        self.res1 = ResidualTemporalBlock1D(in_channels, in_channels // 2, embed_dim=embed_dim)
+        self.down1 = Downsample1D(out_channels // 2, use_conv=True)
+        self.res2 = ResidualTemporalBlock1D(in_channels // 2, in_channels // 4, embed_dim=embed_dim)
+        self.down2 = Downsample1D(out_channels // 4, use_conv=True)
+
+    def forward(self, x, temb=None):
+        x = self.res1(x, temb)
+        x = self.down1(x)
+        x = self.res2(x, temb)
+        x = self.down2(x)
+        return x
+
+
 class MidResTemporalBlock1D(nn.Module):
     def __init__(
         self,
@@ -307,6 +327,8 @@ def get_mid_block(mid_block_type, num_layers, in_channels, out_channels, embed_d
             embed_dim=embed_dim,
             add_downsample=add_downsample,
         )
+    elif mid_block_type == "ValueFunctionMidBlock1D":
+        return ValueFunctionMidBlock1D(in_channels=in_channels, out_channels=out_channels, embed_dim=embed_dim)
     raise ValueError(f"{mid_block_type} does not exist.")
 
 

tests/test_models_unet.py

@@ -22,6 +22,7 @@
 
 from diffusers import UNet1DModel, UNet2DConditionModel, UNet2DModel
 from diffusers.utils import floats_tensor, slow, torch_device
+from regex import subf
 
 from .test_modeling_common import ModelTesterMixin
 
@@ -489,7 +490,7 @@ def prepare_init_args_and_inputs_for_common(self):
 
     def test_from_pretrained_hub(self):
         model, loading_info = UNet1DModel.from_pretrained(
-            "fusing/ddpm-unet-rl-hopper-hor128", output_loading_info=True
+            "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True, subfolder="unet"
         )
         self.assertIsNotNone(model)
         self.assertEqual(len(loading_info["missing_keys"]), 0)
@@ -500,7 +501,7 @@ def test_from_pretrained_hub(self):
         assert image is not None, "Make sure output is not None"
 
     def test_output_pretrained(self):
-        model = UNet1DModel.from_pretrained("fusing/ddpm-unet-rl-hopper-hor128")
+        model = UNet1DModel.from_pretrained("bglick13/hopper-medium-v2-value-function-hor32", subfolder="unet")
         torch.manual_seed(0)
         if torch.cuda.is_available():
             torch.cuda.manual_seed_all(0)
@@ -517,7 +518,8 @@ def test_output_pretrained(self):
 
         output_slice = output[0, -3:, -3:].flatten()
         # fmt: off
-        expected_output_slice = torch.tensor([-0.2714, 0.1042, -0.0794, -0.2820, 0.0803, -0.0811, -0.2345, 0.0580, -0.0584])
+        expected_output_slice = torch.tensor([-2.137172  ,  1.1426016 ,  0.3688687 , -0.766922  ,  0.7303146 ,
+        0.11038864, -0.4760633 ,  0.13270172,  0.02591348])
         # fmt: on
         self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-3))
 
@@ -565,10 +567,10 @@ def prepare_init_args_and_inputs_for_common(self):
 
     def test_from_pretrained_hub(self):
         unet, loading_info = UNet1DModel.from_pretrained(
-            "bglick13/hopper-medium-v2-unet-hor32", output_loading_info=True
+            "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True, subfolder="unet"
         )
         value_function, vf_loading_info = UNet1DModel.from_pretrained(
-            "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True
+            "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True, subfolder="value_function"
         )
         self.assertIsNotNone(unet)
         self.assertEqual(len(loading_info["missing_keys"]), 0)
@@ -583,7 +585,7 @@ def test_from_pretrained_hub(self):
 
     def test_output_pretrained(self):
         value_function, vf_loading_info = UNet1DModel.from_pretrained(
-            "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True
+            "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True, subfolder="value_function"
         )
         torch.manual_seed(0)
         if torch.cuda.is_available():
@@ -600,7 +602,7 @@ def test_output_pretrained(self):
             output = value_function(noise, time_step).sample
 
         # fmt: off
-        expected_output_slice = torch.tensor([207.0272] * seq_len)
+        expected_output_slice = torch.tensor([165.25] * seq_len)
         # fmt: on
         self.assertTrue(torch.allclose(output, expected_output_slice, rtol=1e-3))