Allow converting Flax to PyTorch by adding a "from_flax" keyword #1900

camenduru · 2023-01-03T12:41:52Z

everything changed I am confused 😐 this probably not works @patrickvonplaten please help me

HuggingFaceDocBuilderDev · 2023-01-03T12:53:28Z

The documentation is not available anymore as the PR was closed or merged.

camenduru · 2023-01-03T16:54:52Z

it is working like this

!git clone https://huggingface.co/camenduru/plushies
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("/content/plushies", safety_checker=None, from_flax=True).to("cpu")

but if I do

from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("camenduru/plushies", safety_checker=None, from_flax=True).to("cpu")

it is downloading safetensors 😐

camenduru · 2023-01-03T17:05:45Z

if repo contains safetensors it is downloading safetensors maybe this is not a bug maybe safetensors not implemented in flax yet I will try different model

camenduru · 2023-01-03T17:07:54Z

if only flax model in repo it is skipping 😐

pcuenca · 2023-01-03T21:05:23Z

Hi @camenduru!

This is a great effort, however network downloading is a bit complicated. When you clone the repo, all the files are available locally and you select the flax ones, but downloading attempts to only retrieve the files that will be needed. I haven't had enough time to study your code in depth, but I believe you should manipulate allow_patterns and ignore_patterns here when loading from flax weights. Otherwise, the function snapshot_download would not download them.

For easier debugging, I would concentrate on a repo that only has flax weights and not safetensors. When that works, I'd suggest ignoring all the safetensors logic when loading from flax, to make it simpler.

I hope that helps. I'll try to test it better tomorrow. Thanks a lot for working on this!

camenduru · 2023-01-04T08:19:30Z

@pcuenca thanks ❤ I have a question I found this https://setuptools.pypa.io/en/latest/userguide/development_mode.html @patil-suraj taught me pip install -e ".[dev]" is it possible with this -e ".[dev]" like when I change the code I want to test without uninstalling and installing diffusers

pcuenca · 2023-01-04T08:34:53Z

@pcuenca thanks ❤ I have a question I found this https://setuptools.pypa.io/en/latest/userguide/development_mode.html @patil-suraj taught me pip install -e ".[dev]" is it possible with this -e ".[dev]" like when I change the code I want to test without uninstalling and installing diffusers

Absolutely! If you install with -e then any changes will become live, but you'll only see them in a new Python process. This means that if you are developing in a notebook you have to restart the kernel. If you are using a script or a debugging session in VS Code, you have to kill it and start again.

camenduru · 2023-01-04T09:06:01Z

thanks @pcuenca ❤restart the kernel worked 🎉 also

from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("camenduru/plushies", safety_checker=None, from_flax=True).to("cpu")

working now 🎉🎊✨

camenduru · 2023-01-04T10:34:06Z

how can I pass this test should I edit the test 🤔

=========================== short test summary info ============================
FAILED tests/test_pipelines.py::DownloadTests::test_download_only_pytorch - assert not True
 +  where True = any(<generator object DownloadTests.test_download_only_pytorch.<locals>.<genexpr> at 0x7f270e6e6d60>)
===== 1 failed, 670 passed, 257 skipped, 192 warnings in 163.49s (0:02:43) =====

src/diffusers/pipelines/pipeline_utils.py

camenduru · 2023-01-05T20:41:40Z

from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("camenduru/plushies", safety_checker=None, from_flax=True).to("cpu")
pipe.save_pretrained("pt")

from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("/content/pt", safety_checker=None).to("cuda")
image = pipe("duck", num_inference_steps=50).images[0]
display(image)

from diffusers import FlaxStableDiffusionPipeline
pipe, params = FlaxStableDiffusionPipeline.from_pretrained("camenduru/plushies-pt", from_pt=True)
pipe.save_pretrained("flax", params=params)

from diffusers import FlaxStableDiffusionPipeline
pipe, params = FlaxStableDiffusionPipeline.from_pretrained("/workspaces/flax", dtype=jax.numpy.bfloat16, safety_checker=None)
params = replicate(params)

real_seed = random.randint(0, 2147483647)
prng_seed = jax.random.PRNGKey(real_seed)
prng_seed = jax.random.split(prng_seed, jax.device_count())
num_samples = jax.device_count()
prompt = "duck"
prompt_n = num_samples * [prompt]
prompt_ids = pipe.prepare_inputs(prompt_n)
prompt_ids = shard(prompt_ids)
images = pipe(prompt_ids, params, prng_seed, jit=True).images
images = pipe.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))
display(images[0])

camenduru · 2023-01-05T20:43:53Z

oh no 😐 I did something wrong camenduru requested review from patrickvonplaten and removed request for pcuenca and patil-suraj now

camenduru · 2023-01-05T20:48:04Z

how can I add reviewers back oops sorry

patrickvonplaten

I haven't tried it out, but the changes to modeling_utils.py look good to me! Also all tests are passing 😍

Great job @camenduru - this was one of the harder PRs!

@pcuenca @patil-suraj mind taking a look here as well. Would like to have 2 more 👀 on this one as it touches core functionality.

pcuenca

Looks great, thanks a lot @camenduru! I just left a few minor comments. I'll try to test it tomorrow, but I agree with Patrick that it'd be really cool if we could add a simple test that shows how this works.

Great job!

src/diffusers/models/modeling_pytorch_flax_utils.py

pcuenca · 2023-01-06T23:45:55Z

src/diffusers/models/modeling_pytorch_flax_utils.py

+    else:
+        logger.warning(f"All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n")


This is the ok case, right? I wonder if we should just skip this warning.

pcuenca · 2023-01-06T23:46:22Z

src/diffusers/models/modeling_pytorch_flax_utils.py

+    else:
+        logger.warning(
+            f"All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n"
+            "If your task is similar to the task the model of the checkpoint was trained on, "
+            f"you can already use {pt_model.__class__.__name__} for predictions without further training."
+        )


Same comment here, this means everything went well, doesn't it?

yes deleted

src/diffusers/models/modeling_utils.py

Co-authored-by: Pedro Cuenca <[email protected]>

camenduru · 2023-01-07T10:18:43Z

thanks @pcuenca ❤

if we could add a simple test that shows how this works.

I have a question about test should I write test other then this #1900 (comment)

patrickvonplaten · 2023-01-12T18:59:24Z

Added a test, conversion seems to be working perfectly! Think we can merge this one :-)

Amazing job @camenduru !

camenduru · 2023-01-12T20:50:26Z

woohoo 🥳 thanks @patrickvonplaten ❤ @pcuenca ❤ @patil-suraj ❤ 🥳 🎉 🎊

…gingface#1900) * from_flax * oops * oops * make style with pip install -e ".[dev]" * oops * now code quality happy 😋 * allow_patterns += FLAX_WEIGHTS_NAME * Update src/diffusers/pipelines/pipeline_utils.py Co-authored-by: Patrick von Platen <[email protected]> * Update src/diffusers/pipelines/pipeline_utils.py Co-authored-by: Patrick von Platen <[email protected]> * Update src/diffusers/pipelines/pipeline_utils.py Co-authored-by: Patrick von Platen <[email protected]> * Update src/diffusers/pipelines/pipeline_utils.py Co-authored-by: Patrick von Platen <[email protected]> * Update src/diffusers/models/modeling_utils.py Co-authored-by: Patrick von Platen <[email protected]> * Update src/diffusers/pipelines/pipeline_utils.py Co-authored-by: Patrick von Platen <[email protected]> * for test * bye bye is_flax_available() * oops * Update src/diffusers/models/modeling_pytorch_flax_utils.py Co-authored-by: Pedro Cuenca <[email protected]> * Update src/diffusers/models/modeling_pytorch_flax_utils.py Co-authored-by: Pedro Cuenca <[email protected]> * Update src/diffusers/models/modeling_pytorch_flax_utils.py Co-authored-by: Pedro Cuenca <[email protected]> * Update src/diffusers/models/modeling_utils.py Co-authored-by: Pedro Cuenca <[email protected]> * Update src/diffusers/models/modeling_utils.py Co-authored-by: Pedro Cuenca <[email protected]> * make style * add test * finihs Co-authored-by: Patrick von Platen <[email protected]> Co-authored-by: Pedro Cuenca <[email protected]>

camenduru added 3 commits January 3, 2023 12:35

from_flax

9a5ad09

oops

91f238f

oops

58f3b45

camenduru added 3 commits January 3, 2023 12:56

make style with pip install -e ".[dev]"

f6061ba

Merge branch 'huggingface:main' into from_flax_v3

78dbd4c

oops

75490f1

camenduru changed the title ~~From flax~~ [WIP] From Flax Jan 3, 2023

camenduru added 2 commits January 3, 2023 17:20

now code quality happy 😋

79fd2f3

Merge branch 'huggingface:main' into from_flax_v3

b67334a

patil-suraj requested review from patil-suraj, patrickvonplaten and pcuenca January 3, 2023 17:44

allow_patterns += FLAX_WEIGHTS_NAME

963af41

Merge branch 'huggingface:main' into from_flax_v3

ed36557

camenduru mentioned this pull request Jan 4, 2023

[Question][Solved] How can I use color correction .pt in diffusers #1647

Closed

Merge branch 'main' into from_flax_v3

2d7d37f

camenduru marked this pull request as ready for review January 4, 2023 19:35

patrickvonplaten reviewed Jan 4, 2023

View reviewed changes

src/diffusers/pipelines/pipeline_utils.py Outdated Show resolved Hide resolved

patrickvonplaten reviewed Jan 4, 2023

View reviewed changes

src/diffusers/pipelines/pipeline_utils.py Outdated Show resolved Hide resolved

patrickvonplaten reviewed Jan 4, 2023

View reviewed changes

src/diffusers/pipelines/pipeline_utils.py Show resolved Hide resolved

patrickvonplaten reviewed Jan 4, 2023

View reviewed changes

src/diffusers/pipelines/pipeline_utils.py Show resolved Hide resolved

oops

8d2f05b

camenduru requested review from patrickvonplaten and removed request for patil-suraj and pcuenca January 5, 2023 20:42

Merge branch 'main' into from_flax_v3

fdee0dc

patrickvonplaten approved these changes Jan 5, 2023

View reviewed changes

patrickvonplaten requested review from patil-suraj and pcuenca January 5, 2023 22:18

pcuenca approved these changes Jan 6, 2023

View reviewed changes

camenduru and others added 7 commits January 7, 2023 12:59

Update src/diffusers/models/modeling_pytorch_flax_utils.py

3f4fdc5

Co-authored-by: Pedro Cuenca <[email protected]>

Update src/diffusers/models/modeling_pytorch_flax_utils.py

d132884

Co-authored-by: Pedro Cuenca <[email protected]>

Update src/diffusers/models/modeling_pytorch_flax_utils.py

ea48d19

Co-authored-by: Pedro Cuenca <[email protected]>

Update src/diffusers/models/modeling_utils.py

dfe1b3e

Co-authored-by: Pedro Cuenca <[email protected]>

Update src/diffusers/models/modeling_utils.py

b05232b

Co-authored-by: Pedro Cuenca <[email protected]>

Merge branch 'main' into from_flax_v3

7abea4f

make style

7599a23

camenduru and others added 2 commits January 10, 2023 17:09

Merge branch 'huggingface:main' into from_flax_v3

4eb3a9c

add test

4ff9dc0

patrickvonplaten changed the title ~~[WIP] From Flax~~ From Flax Jan 12, 2023

finihs

990d43d

patrickvonplaten changed the title ~~From Flax~~ Allow converting Flax to PyTorch by adding a "from_flax" keyword Jan 12, 2023

patrickvonplaten merged commit f73ed17 into huggingface:main Jan 12, 2023

Lime-Cakes mentioned this pull request Jan 20, 2023

Allow Flax model to save as pytorch (If converted from pytorch) #1739

Closed

		else:
		logger.warning(f"All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n")

Allow converting Flax to PyTorch by adding a "from_flax" keyword #1900

Allow converting Flax to PyTorch by adding a "from_flax" keyword #1900

Uh oh!

Conversation

camenduru commented Jan 3, 2023

Uh oh!

HuggingFaceDocBuilderDev commented Jan 3, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

camenduru commented Jan 3, 2023

Uh oh!

camenduru commented Jan 3, 2023

Uh oh!

camenduru commented Jan 3, 2023

Uh oh!

pcuenca commented Jan 3, 2023

Uh oh!

camenduru commented Jan 4, 2023

Uh oh!

pcuenca commented Jan 4, 2023

Uh oh!

camenduru commented Jan 4, 2023

Uh oh!

camenduru commented Jan 4, 2023

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

camenduru commented Jan 5, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

camenduru commented Jan 5, 2023

Uh oh!

camenduru commented Jan 5, 2023

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

pcuenca left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pcuenca Jan 6, 2023

Choose a reason for hiding this comment

Uh oh!

camenduru Jan 7, 2023

Choose a reason for hiding this comment

Uh oh!

pcuenca Jan 6, 2023

Choose a reason for hiding this comment

Uh oh!

camenduru Jan 7, 2023

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

camenduru commented Jan 7, 2023

Uh oh!

patrickvonplaten commented Jan 12, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

camenduru commented Jan 12, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

HuggingFaceDocBuilderDev commented Jan 3, 2023 •

edited

Loading

camenduru commented Jan 5, 2023 •

edited

Loading

patrickvonplaten commented Jan 12, 2023 •

edited

Loading