Add Flax stable diffusion img2img pipeline #1355

dhruvrnaik · 2022-11-21T15:50:44Z

Added Flax SD pipeline for Image 2 Image

TODO:

Update README
Fix generation bug

The pipeline is working, but the result is not as expected.

Input image:
Prompt = "A fantasy landscape, trending on artstation"
Outputs:

Params:
strength=0.75,
num_inference_steps=50,
guidance_scale = 7.5

How to run this:

import requests
from io import BytesIO
from PIL import Image

url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"

response = requests.get(url)
init_img = Image.open(BytesIO(response.content)).convert("RGB")
init_img = init_img.resize((768, 512))
prompts = "A fantasy landscape, trending on artstation"
dtype=jnp.bfloat16
pipeline, params = FlaxStableDiffusionImg2ImgPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", revision="flax",
    dtype=dtype,
)
def create_key(seed=0):
    return jax.random.PRNGKey(seed)
rng = create_key(0)
rng = jax.random.split(rng, jax.device_count())

prompt_ids, imgs = pipeline.prepare_inputs(prompt=[prompts]*jax.device_count(), init_image = [init_img]*jax.device_count())
p_params = replicate(params)
prompt_ids = shard(prompt_ids)
imgs = shard(imgs)

output = pipeline(
    prompt_ids=prompt_ids, 
    init_images=imgs, 
    params=p_params, 
    prng_seed=rng, 
    strength=0.75, 
    num_inference_steps=50, 
    jit=True, 
    guidance_scale=7.5,
height=512,width=768).images

HuggingFaceDocBuilderDev · 2022-11-21T15:55:13Z

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

dhruvrnaik · 2022-11-22T15:17:46Z

cc: @patil-suraj if you or someone else could review this, It'd be great. Trying to fix the image result, but a review can help fix any other issues.

dhruvrnaik · 2022-11-25T03:12:52Z

Fixed it. Forgot to update the for loop with the right timesteps.

patrickvonplaten · 2022-11-30T12:01:16Z

Cool ! @dhruvrnaik do you need help finishing this PR?

dhruvrnaik · 2022-12-02T01:09:26Z

@patrickvonplaten I am mostly done with this. I just need to update the docs, which I will finish tonight. I have tested this on a TPU v3-8 pod and it works as expected. Ready for a review, otherwise.
I will fix the black/isort test fails too.

pcuenca · 2022-12-02T11:27:54Z

@patrickvonplaten I am mostly done with this. I just need to update the docs, which I will finish tonight. I have tested this on a TPU v3-8 pod and it works as expected. Ready for a review, otherwise.
I will fix the black/isort test fails too.

Excellent! I'll check it out too. Please, let me know when you are done with your changes :)
Thanks a lot!

pcuenca

Looks good! However, it still doesn't work for me. We also need to expose the class name up through the __init__.py chain so it's visible in diffusers.

Please, make sure you push any other changes you may have locally, and let us know if you need help driving this home :)

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

README.md

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py

README.md

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py

pcuenca · 2022-12-03T20:41:20Z

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py

+        if debug:
+            # run with python for loop
+            for i in range(len(timesteps)):
+                latents, timesteps, scheduler_state = loop_body(i, (latents, timesteps, scheduler_state))
+        else:
+            latents, _, _ = jax.lax.fori_loop(0, len(timesteps), loop_body, (latents, timesteps, scheduler_state))


I'm not sure about this. Wouldn't it be possible to start the loop at the initial effective timestep (instead of 0), and avoid having to pass it around?

This is also done. Please have a look when you can

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py

…diffusion_img2img.py Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

patrickvonplaten · 2022-12-15T20:08:09Z

Hey @dhruvrnaik,

Thanks a lot for the PR! Could you maybe run make fix-copies once and correct the:

 ❱  45 class FlaxStableDiffusionImg2ImgPipeline(FlaxDiffusionPipeline):       │
│    46 │   r"""                                                               │
│    47 │   Pipeline for image-to-image generation using Stable Diffusion.     │
│    48                                                                        │
│                                                                              │
│ /usr/local/lib/python3.8/site-packages/diffusers/pipelines/stable_diffusion/ │
│ pipeline_flax_stable_diffusion_img2img.py:270 in                             │
│ FlaxStableDiffusionImg2ImgPipeline                                           │
│                                                                              │
│   267 │   │   prompt_ids: jnp.array,                                         │
│   268 │   │   image: jnp.array,                                              │
│   269 │   │   params: Union[Dict, FrozenDict],                               │
│ ❱ 270 │   │   prng_seed: Union[jax.random.KeyArray, jax.Array],              │
│   271 │   │   num_inference_steps: int = 50,                                 │
│   272 │   │   height: int = 512,                                             │
│   273 │   │   width: int = 512,                                              │
╰──────────────────────────────────────────────────────────────────────────────╯
AttributeError: module 'jax' has no attribute 'Array'

error message of the PR Documentation to get this merged? :-)

dhruvrnaik · 2022-12-19T05:01:42Z

Should be good to go @patrickvonplaten

pcuenca · 2022-12-19T10:20:44Z

@dhruvrnaik I'll test it today.

patrickvonplaten · 2022-12-20T00:15:13Z

@pcuenca, feel free to merge whenever :-)

pcuenca · 2022-12-20T15:23:19Z

Thanks a lot @dhruvrnaik!

* add flax img2img pipeline * update pipeline * black format file * remove argg from get_timesteps * update get_timesteps * fix bug: make use of timesteps for for_loop * black file * black, isort, flake8 * update docstring * update readme * update flax img2img readme * update sd pipeline init * Update src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * Update src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * update inits * revert change * update var name to image, typo * update readme * return new t_start instead of modified timestep * black format * isort files * update docs * fix-copies * update prng_seed typing Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

skirsten · 2022-12-21T16:33:47Z

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py

+
+        t_start = self.get_timestep_start(num_inference_steps, strength, scheduler_state)
+        latent_timestep = scheduler_state.timesteps[t_start : t_start + 1].repeat(batch_size)
+        init_latents = self.scheduler.add_noise(init_latents, noise, latent_timestep)


As of #1661, the add_noise function now requires the scheduler params:

init_latents = self.scheduler.add_noise(params["scheduler"], init_latents, noise, latent_timestep)

Thanks @skirsten! Would you like to open a PR for that?

In addition, we should create some tests for this new pipeline too.

I can work on the tests. Will check out the PR too

Here is the PR that also adds the params to add_noise: #1824

* add flax img2img pipeline * update pipeline * black format file * remove argg from get_timesteps * update get_timesteps * fix bug: make use of timesteps for for_loop * black file * black, isort, flake8 * update docstring * update readme * update flax img2img readme * update sd pipeline init * Update src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * Update src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * update inits * revert change * update var name to image, typo * update readme * return new t_start instead of modified timestep * black format * isort files * update docs * fix-copies * update prng_seed typing Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

add flax img2img pipeline

b3f6343

dhruvrnaik marked this pull request as draft November 21, 2022 15:56

dhruvrnaik added 3 commits November 21, 2022 11:09

update pipeline

e637e83

Merge branch 'main' into flax-img2img

30293cc

Merge branch 'main' into flax-img2img

20b6ce3

dhruvrnaik added 3 commits November 22, 2022 10:24

black format file

e80efe9

remove argg from get_timesteps

9603d75

update get_timesteps

a1b27ef

patil-suraj assigned patil-suraj and pcuenca Nov 23, 2022

dhruvrnaik added 2 commits November 24, 2022 21:59

fix bug: make use of timesteps for for_loop

7122386

Merge branch 'main' into flax-img2img

72fdb95

dhruvrnaik marked this pull request as ready for review November 25, 2022 03:09

dhruvrnaik added 3 commits November 25, 2022 12:11

black file

727fa1d

black, isort, flake8

e7d6687

update docstring

c8787c8

dhruvrnaik added 3 commits December 2, 2022 22:29

update readme

94a3e93

update flax img2img readme

7696bf7

update sd pipeline init

c5a4275

pcuenca reviewed Dec 3, 2022

View reviewed changes

src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py Outdated Show resolved Hide resolved

dhruvrnaik and others added 3 commits December 4, 2022 21:18

Update src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_…

83f2f77

…diffusion_img2img.py Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

Update src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_…

bc4abd6

…diffusion_img2img.py Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

update inits

ff2be25

dhruvrnaik added 4 commits December 6, 2022 13:39

revert change

6cec0b8

update var name to image, typo

e124ed9

update readme

c1e6996

return new t_start instead of modified timestep

932a74e

dhruvrnaik requested a review from pcuenca December 12, 2022 00:34

dhruvrnaik added 4 commits December 11, 2022 19:35

Merge branch 'main' into flax-img2img

6602ec3

black format

e8dff83

isort files

d06f8c3

update docs

b9fa9b7

dhruvrnaik added 2 commits December 19, 2022 04:37

fix-copies

fb0af34

update prng_seed typing

cb88c46

pcuenca approved these changes Dec 20, 2022

View reviewed changes

pcuenca merged commit a9190ba into huggingface:main Dec 20, 2022

skirsten reviewed Dec 21, 2022

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Flax stable diffusion img2img pipeline #1355

Add Flax stable diffusion img2img pipeline #1355

dhruvrnaik commented Nov 21, 2022 •

edited

HuggingFaceDocBuilderDev commented Nov 21, 2022 •

edited

dhruvrnaik commented Nov 22, 2022

dhruvrnaik commented Nov 25, 2022 •

edited

patrickvonplaten commented Nov 30, 2022

dhruvrnaik commented Dec 2, 2022

pcuenca commented Dec 2, 2022

pcuenca left a comment

pcuenca Dec 3, 2022

dhruvrnaik Dec 12, 2022

patrickvonplaten commented Dec 15, 2022

dhruvrnaik commented Dec 19, 2022

pcuenca commented Dec 19, 2022 •

edited

patrickvonplaten commented Dec 20, 2022

pcuenca commented Dec 20, 2022

skirsten Dec 21, 2022

pcuenca Dec 22, 2022

dhruvrnaik Dec 22, 2022

skirsten Dec 25, 2022

Add Flax stable diffusion img2img pipeline #1355

Add Flax stable diffusion img2img pipeline #1355

Conversation

dhruvrnaik commented Nov 21, 2022 • edited

HuggingFaceDocBuilderDev commented Nov 21, 2022 • edited

dhruvrnaik commented Nov 22, 2022

dhruvrnaik commented Nov 25, 2022 • edited

patrickvonplaten commented Nov 30, 2022

dhruvrnaik commented Dec 2, 2022

pcuenca commented Dec 2, 2022

pcuenca left a comment

Choose a reason for hiding this comment

pcuenca Dec 3, 2022

Choose a reason for hiding this comment

dhruvrnaik Dec 12, 2022

Choose a reason for hiding this comment

patrickvonplaten commented Dec 15, 2022

dhruvrnaik commented Dec 19, 2022

pcuenca commented Dec 19, 2022 • edited

patrickvonplaten commented Dec 20, 2022

pcuenca commented Dec 20, 2022

skirsten Dec 21, 2022

Choose a reason for hiding this comment

pcuenca Dec 22, 2022

Choose a reason for hiding this comment

dhruvrnaik Dec 22, 2022

Choose a reason for hiding this comment

skirsten Dec 25, 2022

Choose a reason for hiding this comment

dhruvrnaik commented Nov 21, 2022 •

edited

HuggingFaceDocBuilderDev commented Nov 21, 2022 •

edited

dhruvrnaik commented Nov 25, 2022 •

edited

pcuenca commented Dec 19, 2022 •

edited