Merge branch 'main' into overhaul-release-workflow

.github/workflows/claude_review.yml

@@ -32,9 +32,6 @@ jobs:
       )
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
-        with:
-          fetch-depth: 1
       - uses: anthropics/claude-code-action@v1
         with:
           anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}

.github/workflows/pypi_publish.yaml

@@ -1,73 +1,45 @@
-# Adapted from https://blog.deepjyoti30.dev/pypi-release-github-action
-
 name: PyPI release
 
 on:
   workflow_dispatch:
   push:
     tags:
-      - "*"
+      - "v*"
 
 jobs:
-  find-and-checkout-latest-branch:
+  build-and-test:
     runs-on: ubuntu-22.04
-    outputs:
-      latest_branch: ${{ steps.set_latest_branch.outputs.latest_branch }}
     steps:
-      - name: Checkout Repo
+      - name: Checkout repo
         uses: actions/checkout@v6
 
       - name: Set up Python
         uses: actions/setup-python@v6
         with:
-          python-version: '3.10'
+          python-version: "3.10"
 
-      - name: Fetch latest branch
-        id: fetch_latest_branch
+      - name: Fetch and checkout latest release branch
         run: |
           pip install -U requests packaging
           LATEST_BRANCH=$(python utils/fetch_latest_release_branch.py)
           echo "Latest branch: $LATEST_BRANCH"
-          echo "latest_branch=$LATEST_BRANCH" >> $GITHUB_ENV
+          git fetch origin "$LATEST_BRANCH"
+          git checkout "$LATEST_BRANCH"
 
-      - name: Set latest branch output
-        id: set_latest_branch
-        run: echo "::set-output name=latest_branch::${{ env.latest_branch }}"
-
-  release:
-    needs: find-and-checkout-latest-branch
-    runs-on: ubuntu-22.04
-
-    steps:
-      - name: Checkout Repo
-        uses: actions/checkout@v6
-        with:
-          ref: ${{ needs.find-and-checkout-latest-branch.outputs.latest_branch }}
-
-      - name: Setup Python
-        uses: actions/setup-python@v6
-        with:
-          python-version: "3.10"
-
-      - name: Install dependencies
+      - name: Install build dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -U setuptools wheel twine
+          pip install -U build
           pip install -U torch --index-url https://download.pytorch.org/whl/cpu
 
       - name: Build the dist files
-        run: python setup.py bdist_wheel && python setup.py sdist
+        run: python -m build
 
-      - name: Publish to the test PyPI
-        env:
-          TWINE_USERNAME: ${{ secrets.TEST_PYPI_USERNAME }}
-          TWINE_PASSWORD: ${{ secrets.TEST_PYPI_PASSWORD }}
-        run: twine upload dist/* -r pypitest --repository-url=https://test.pypi.org/legacy/
+      - name: Install from built wheel
+        run: pip install dist/*.whl
 
       - name: Test installing diffusers and importing
         run: |
-          pip install diffusers && pip uninstall diffusers -y
-          pip install -i https://test.pypi.org/simple/ diffusers
           pip install -U transformers
           python utils/print_env.py
           python -c "from diffusers import __version__; print(__version__)"
@@ -75,8 +47,26 @@ jobs:
           python -c "from diffusers import DiffusionPipeline; pipe = DiffusionPipeline.from_pretrained('hf-internal-testing/tiny-stable-diffusion-pipe', safety_checker=None); pipe('ah suh du')"
           python -c "from diffusers import *"
 
+      - name: Upload build artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-dist
+          path: dist/
+
+  publish-to-pypi:
+    needs: build-and-test
+    if: startsWith(github.ref, 'refs/tags/')
+    runs-on: ubuntu-22.04
+    environment: pypi-release
+    permissions:
+      id-token: write
+
+    steps:
+      - name: Download build artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: python-dist
+          path: dist/
+
       - name: Publish to PyPI
-        env:
-          TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
-          TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
-        run: twine upload dist/* -r pypi
+        uses: pypa/gh-action-pypi-publish@release/v1

tests/models/transformers/test_models_transformer_bria.py

@@ -13,31 +13,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
+import unittest
 
 import torch
 
 from diffusers import BriaTransformer2DModel
 from diffusers.models.attention_processor import FluxIPAdapterJointAttnProcessor2_0
 from diffusers.models.embeddings import ImageProjection
-from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
-from ..testing_utils import (
-    BaseModelTesterConfig,
-    IPAdapterTesterMixin,
-    LoraHotSwappingForModelTesterMixin,
-    LoraTesterMixin,
-    ModelTesterMixin,
-    TorchCompileTesterMixin,
-    TrainingTesterMixin,
-)
+from ..test_modeling_common import LoraHotSwappingForModelTesterMixin, ModelTesterMixin, TorchCompileTesterMixin
 
 
 enable_full_determinism()
 
 
-def create_bria_ip_adapter_state_dict(model) -> dict[str, dict[str, Any]]:
+def create_bria_ip_adapter_state_dict(model):
+    # "ip_adapter" (cross-attention weights)
     ip_cross_attn_state_dict = {}
     key_id = 0
 
@@ -58,8 +50,11 @@ def create_bria_ip_adapter_state_dict(model) -> dict[str, dict[str, Any]]:
                 f"{key_id}.to_v_ip.bias": sd["to_v_ip.0.bias"],
             }
         )
+
         key_id += 1
 
+    # "image_proj" (ImageProjection layer weights)
+
     image_projection = ImageProjection(
         cross_attention_dim=model.config["joint_attention_dim"],
         image_embed_dim=model.config["pooled_projection_dim"],
@@ -78,36 +73,53 @@ def create_bria_ip_adapter_state_dict(model) -> dict[str, dict[str, Any]]:
     )
 
     del sd
-    return {"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict}
+    ip_state_dict = {}
+    ip_state_dict.update({"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict})
+    return ip_state_dict
 
 
-class BriaTransformerTesterConfig(BaseModelTesterConfig):
-    @property
-    def model_class(self):
-        return BriaTransformer2DModel
+class BriaTransformerTests(ModelTesterMixin, unittest.TestCase):
+    model_class = BriaTransformer2DModel
+    main_input_name = "hidden_states"
+    # We override the items here because the transformer under consideration is small.
+    model_split_percents = [0.8, 0.7, 0.7]
+
+    # Skip setting testing with default: AttnProcessor
+    uses_custom_attn_processor = True
 
     @property
-    def main_input_name(self) -> str:
-        return "hidden_states"
+    def dummy_input(self):
+        batch_size = 1
+        num_latent_channels = 4
+        num_image_channels = 3
+        height = width = 4
+        sequence_length = 48
+        embedding_dim = 32
 
-    @property
-    def model_split_percents(self) -> list:
-        return [0.8, 0.7, 0.7]
+        hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(torch_device)
+        encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
+        text_ids = torch.randn((sequence_length, num_image_channels)).to(torch_device)
+        image_ids = torch.randn((height * width, num_image_channels)).to(torch_device)
+        timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size)
 
-    @property
-    def output_shape(self) -> tuple:
-        return (16, 4)
-
-    @property
-    def input_shape(self) -> tuple:
-        return (16, 4)
-
-    @property
-    def generator(self):
-        return torch.Generator("cpu").manual_seed(0)
-
-    def get_init_dict(self) -> dict:
         return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "img_ids": image_ids,
+            "txt_ids": text_ids,
+            "timestep": timestep,
+        }
+
+    @property
+    def input_shape(self):
+        return (16, 4)
+
+    @property
+    def output_shape(self):
+        return (16, 4)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
             "patch_size": 1,
             "in_channels": 4,
             "num_layers": 1,
@@ -119,35 +131,11 @@ class BriaTransformerTesterConfig(BaseModelTesterConfig):
             "axes_dims_rope": [0, 4, 4],
         }
 
-    def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]:
-        num_latent_channels = 4
-        num_image_channels = 3
-        height = width = 4
-        sequence_length = 48
-        embedding_dim = 32
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
 
-        return {
-            "hidden_states": randn_tensor(
-                (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device
-            ),
-            "encoder_hidden_states": randn_tensor(
-                (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
-            ),
-            "img_ids": randn_tensor(
-                (height * width, num_image_channels), generator=self.generator, device=torch_device
-            ),
-            "txt_ids": randn_tensor(
-                (sequence_length, num_image_channels), generator=self.generator, device=torch_device
-            ),
-            "timestep": torch.tensor([1.0]).to(torch_device).expand(batch_size),
-        }
-
-
-class TestBriaTransformer(BriaTransformerTesterConfig, ModelTesterMixin):
     def test_deprecated_inputs_img_txt_ids_3d(self):
-        init_dict = self.get_init_dict()
-        inputs_dict = self.get_dummy_inputs()
-
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
         model = self.model_class(**init_dict)
         model.to(torch_device)
         model.eval()
@@ -155,6 +143,7 @@ class TestBriaTransformer(BriaTransformerTesterConfig, ModelTesterMixin):
         with torch.no_grad():
             output_1 = model(**inputs_dict).to_tuple()[0]
 
+        # update inputs_dict with txt_ids and img_ids as 3d tensors (deprecated)
         text_ids_3d = inputs_dict["txt_ids"].unsqueeze(0)
         image_ids_3d = inputs_dict["img_ids"].unsqueeze(0)
 
@@ -167,59 +156,26 @@ class TestBriaTransformer(BriaTransformerTesterConfig, ModelTesterMixin):
         with torch.no_grad():
             output_2 = model(**inputs_dict).to_tuple()[0]
 
-        assert output_1.shape == output_2.shape
-        assert torch.allclose(output_1, output_2, atol=1e-5), (
-            "output with deprecated inputs (img_ids and txt_ids as 3d torch tensors) "
-            "are not equal as them as 2d inputs"
+        self.assertEqual(output_1.shape, output_2.shape)
+        self.assertTrue(
+            torch.allclose(output_1, output_2, atol=1e-5),
+            msg="output with deprecated inputs (img_ids and txt_ids as 3d torch tensors) are not equal as them as 2d inputs",
         )
 
-
-class TestBriaTransformerTraining(BriaTransformerTesterConfig, TrainingTesterMixin):
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {"BriaTransformer2DModel"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
 
-class TestBriaTransformerCompile(BriaTransformerTesterConfig, TorchCompileTesterMixin):
-    pass
+class BriaTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
+    model_class = BriaTransformer2DModel
+
+    def prepare_init_args_and_inputs_for_common(self):
+        return BriaTransformerTests().prepare_init_args_and_inputs_for_common()
 
 
-class TestBriaTransformerIPAdapter(BriaTransformerTesterConfig, IPAdapterTesterMixin):
-    @property
-    def ip_adapter_processor_cls(self):
-        return FluxIPAdapterJointAttnProcessor2_0
+class BriaTransformerLoRAHotSwapTests(LoraHotSwappingForModelTesterMixin, unittest.TestCase):
+    model_class = BriaTransformer2DModel
 
-    def modify_inputs_for_ip_adapter(self, model, inputs_dict):
-        torch.manual_seed(0)
-        cross_attention_dim = getattr(model.config, "joint_attention_dim", 32)
-        image_embeds = torch.randn(1, 1, cross_attention_dim).to(torch_device)
-        inputs_dict.update({"joint_attention_kwargs": {"ip_adapter_image_embeds": image_embeds}})
-        return inputs_dict
-
-    def create_ip_adapter_state_dict(self, model: Any) -> dict[str, dict[str, Any]]:
-        return create_bria_ip_adapter_state_dict(model)
-
-
-class TestBriaTransformerLoRA(BriaTransformerTesterConfig, LoraTesterMixin):
-    pass
-
-
-class TestBriaTransformerLoRAHotSwap(BriaTransformerTesterConfig, LoraHotSwappingForModelTesterMixin):
-    @property
-    def different_shapes_for_compilation(self):
-        return [(4, 4), (4, 8), (8, 8)]
-
-    def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:
-        batch_size = 1
-        num_latent_channels = 4
-        num_image_channels = 3
-        sequence_length = 24
-        embedding_dim = 32
-
-        return {
-            "hidden_states": randn_tensor((batch_size, height * width, num_latent_channels), device=torch_device),
-            "encoder_hidden_states": randn_tensor((batch_size, sequence_length, embedding_dim), device=torch_device),
-            "img_ids": randn_tensor((height * width, num_image_channels), device=torch_device),
-            "txt_ids": randn_tensor((sequence_length, num_image_channels), device=torch_device),
-            "timestep": torch.tensor([1.0]).to(torch_device).expand(batch_size),
-        }
+    def prepare_init_args_and_inputs_for_common(self):
+        return BriaTransformerTests().prepare_init_args_and_inputs_for_common()

tests/models/transformers/test_models_transformer_bria_fibo.py

@@ -13,50 +13,62 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import unittest
+
 import torch
 
 from diffusers import BriaFiboTransformer2DModel
-from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
-from ..testing_utils import (
-    BaseModelTesterConfig,
-    ModelTesterMixin,
-    TorchCompileTesterMixin,
-    TrainingTesterMixin,
-)
+from ..test_modeling_common import ModelTesterMixin
 
 
 enable_full_determinism()
 
 
-class BriaFiboTransformerTesterConfig(BaseModelTesterConfig):
-    @property
-    def model_class(self):
-        return BriaFiboTransformer2DModel
+class BriaFiboTransformerTests(ModelTesterMixin, unittest.TestCase):
+    model_class = BriaFiboTransformer2DModel
+    main_input_name = "hidden_states"
+    # We override the items here because the transformer under consideration is small.
+    model_split_percents = [0.8, 0.7, 0.7]
+
+    # Skip setting testing with default: AttnProcessor
+    uses_custom_attn_processor = True
 
     @property
-    def main_input_name(self) -> str:
-        return "hidden_states"
+    def dummy_input(self):
+        batch_size = 1
+        num_latent_channels = 48
+        num_image_channels = 3
+        height = width = 16
+        sequence_length = 32
+        embedding_dim = 64
+
+        hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(torch_device)
+        encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
+        text_ids = torch.randn((sequence_length, num_image_channels)).to(torch_device)
+        image_ids = torch.randn((height * width, num_image_channels)).to(torch_device)
+        timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size)
+
+        return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "img_ids": image_ids,
+            "txt_ids": text_ids,
+            "timestep": timestep,
+            "text_encoder_layers": [encoder_hidden_states[:, :, :32], encoder_hidden_states[:, :, :32]],
+        }
 
     @property
-    def model_split_percents(self) -> list:
-        return [0.8, 0.7, 0.7]
-
-    @property
-    def output_shape(self) -> tuple:
-        return (256, 48)
-
-    @property
-    def input_shape(self) -> tuple:
+    def input_shape(self):
         return (16, 16)
 
     @property
-    def generator(self):
-        return torch.Generator("cpu").manual_seed(0)
+    def output_shape(self):
+        return (256, 48)
 
-    def get_init_dict(self) -> dict:
-        return {
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
             "patch_size": 1,
             "in_channels": 48,
             "num_layers": 1,
@@ -69,41 +81,9 @@ class BriaFiboTransformerTesterConfig(BaseModelTesterConfig):
             "axes_dims_rope": [0, 4, 4],
         }
 
-    def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]:
-        num_latent_channels = 48
-        num_image_channels = 3
-        height = width = 16
-        sequence_length = 32
-        embedding_dim = 64
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
 
-        encoder_hidden_states = randn_tensor(
-            (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
-        )
-        return {
-            "hidden_states": randn_tensor(
-                (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device
-            ),
-            "encoder_hidden_states": encoder_hidden_states,
-            "img_ids": randn_tensor(
-                (height * width, num_image_channels), generator=self.generator, device=torch_device
-            ),
-            "txt_ids": randn_tensor(
-                (sequence_length, num_image_channels), generator=self.generator, device=torch_device
-            ),
-            "timestep": torch.tensor([1.0]).to(torch_device).expand(batch_size),
-            "text_encoder_layers": [encoder_hidden_states[:, :, :32], encoder_hidden_states[:, :, :32]],
-        }
-
-
-class TestBriaFiboTransformer(BriaFiboTransformerTesterConfig, ModelTesterMixin):
-    pass
-
-
-class TestBriaFiboTransformerTraining(BriaFiboTransformerTesterConfig, TrainingTesterMixin):
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {"BriaFiboTransformer2DModel"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
-
-
-class TestBriaFiboTransformerCompile(BriaFiboTransformerTesterConfig, TorchCompileTesterMixin):
-    pass