[FLAX-PT conversion] test

[mishig25]

Snippet from @patil-suraj that should become test for testing from_pt=True conversion

from diffusers import UNet2DConditionModel, FlaxUNet2DConditionModel
import tempfile
import torch
import numpy as np
import jax.numpy as jnp

model = UNet2DConditionModel(
    block_out_channels=(32, 64),
    layers_per_block=2,
    sample_size=32,
    in_channels=4,
    out_channels=4,
    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
    up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
    cross_attention_dim=32,
)

model = model.eval()

with tempfile.TemporaryDirectory() as tmpdirname:
    model.save_pretrained(tmpdirname)
    flax_model, flax_params = FlaxUNet2DConditionModel.from_pretrained(tmpdirname, from_pt=True)

sample = torch.rand(1, 4, model.config.sample_size, model.config.sample_size)
time = 1
text_emb = torch.rand(1, 77, model.config.sample_size)


# Step 1: Infer with the PT model
torch_output = model(sample, time, text_emb).sample


# Step 2: Infer with JAX model
flax_sample = jnp.array(sample.numpy())
flax_text_emb = jnp.array(text_emb.numpy())

flax_output = flax_model.apply({"params":flax_params}, flax_sample, time, flax_text_emb).sample

# Step 3: Check that the values are close
converted_flax_output = torch.from_numpy(np.array(flax_output))

torch.testing.assert_allclose(converted_flax_output, torch_output, rtol=4e-03, atol=4e-03)

cc: @patil-suraj @younesbelkada

[younesbelkada]

Since #594 has been merged the test should now pass with

torch.testing.assert_allclose(converted_flax_output, torch_output, rtol=1e-04, atol=1e-04)

[patrickvonplaten]

Thanks for putting it here!

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.