fixed a dtype bfloat16 bug in torch_utils.py

[zhangp365]

when generating 1024*1024 image with bfloat16 dtype, there is an exception:
File "/opt/conda/lib/python3.10/site-packages/diffusers/utils/torch_utils.py", line 107, in fourier_filter
x_freq = fftn(x, dim=(-2, -1))
RuntimeError: Unsupported dtype BFloat16

What does this PR do?

fix a bug.

@hlky

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[hlky]

@zhangp365 Thanks! this is when using freeu? Can we keep the check for non-power of 2 images and add another for bfloat16? I think that makes it clearer why we're casting to float32.

[zhangp365]

@hlky Yes, this uses FreeU and sets the pipeline dtype to bfloat16. In this case, when the image is not a non-power-of-2 size, the pipeline runs successfully. However, the standard size image fails to run. Therefore, I believe casting the type to float32 is a safe operation, making the code more robust.

[hlky]

    # Non-power of 2 images must be float32
    if (W & (W - 1)) != 0 or (H & (H - 1)) != 0:
        x = x.to(dtype=torch.float32)
    # fftn does not support bfloat16
    elif x.dtype == torch.bfloat16:
        x = x.to(dtype=torch.float32)

If we always cast someone looking at the function in the future may wonder why. cc @sayakpaul @DN6 WDYT?

[sayakpaul]

Makes sense to me!

[hlky]

@zhangp365 Can you run make style?

[zhangp365]

@zhangp365 Can you run make style?

I tried make style, there are a lot of error:

examples/community/adaptive_mask_inpainting.py:1412:12: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
     |
1411 |     def vis_seg_on_img(self, image, mask):
1412 |         if type(mask) == np.ndarray:
     |            ^^^^^^^^^^^^^^^^^^^^^^^^ E721
1413 |             mask = torch.tensor(mask)
1414 |         v = Visualizer(image, self.coco_metadata, scale=0.5, instance_mode=ColorMode.IMAGE_BW)
     |

examples/research_projects/geodiff/geodiff_molecule_conformation.ipynb:cell 53:59:7: F821 Undefined name `pickle`
   |
58 |   with open(save_path, 'wb') as f:
59 |       pickle.dump(results, f)
   |       ^^^^^^ F821
   |

examples/research_projects/gligen/demo.ipynb:cell 5:13:1: E402 Module level import not at top of cell
   |
11 | gen_boxes = [('a steam boat', [232, 225, 257, 149]), ('a jumping pink dolphin', [21, 249, 189, 123])]
12 |
13 | import numpy as np
   | ^^^^^^^^^^^^^^^^^^ E402
   |

src/diffusers/configuration_utils.py:691:16: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
689 |             if field.name in self._flax_internal_args:
690 |                 continue
691 |             if type(field.default) == dataclasses._MISSING_TYPE:
    |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
692 |                 default_kwargs[field.name] = None
693 |             else:
    |

tests/models/test_modeling_common.py:398:20: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
397 |         model.set_default_attn_processor()
398 |         assert all(type(proc) == AttnProcessorNPU for proc in model.attn_processors.values())
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
399 |         with torch.no_grad():
400 |             if self.forward_requires_fresh_args:
    |

tests/models/test_modeling_common.py:406:20: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
405 |         model.enable_npu_flash_attention()
406 |         assert all(type(proc) == AttnProcessorNPU for proc in model.attn_processors.values())
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
407 |         with torch.no_grad():
408 |             if self.forward_requires_fresh_args:
    |

tests/models/test_modeling_common.py:414:20: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
413 |         model.set_attn_processor(AttnProcessorNPU())
414 |         assert all(type(proc) == AttnProcessorNPU for proc in model.attn_processors.values())
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
415 |         with torch.no_grad():
416 |             if self.forward_requires_fresh_args:
    |

tests/models/test_modeling_common.py:449:20: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
448 |         model.set_default_attn_processor()
449 |         assert all(type(proc) == AttnProcessor for proc in model.attn_processors.values())
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
450 |         with torch.no_grad():
451 |             if self.forward_requires_fresh_args:
    |

tests/models/test_modeling_common.py:457:20: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
456 |         model.enable_xformers_memory_efficient_attention()
457 |         assert all(type(proc) == XFormersAttnProcessor for proc in model.attn_processors.values())
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
458 |         with torch.no_grad():
459 |             if self.forward_requires_fresh_args:
    |

tests/models/test_modeling_common.py:465:20: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
464 |         model.set_attn_processor(XFormersAttnProcessor())
465 |         assert all(type(proc) == XFormersAttnProcessor for proc in model.attn_processors.values())
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
466 |         with torch.no_grad():
467 |             if self.forward_requires_fresh_args:
    |

tests/models/test_modeling_common.py:496:20: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
494 |             return
495 |
496 |         assert all(type(proc) == AttnProcessor2_0 for proc in model.attn_processors.values())
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
497 |         with torch.no_grad():
498 |             if self.forward_requires_fresh_args:
    |

tests/models/test_modeling_common.py:504:20: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
503 |         model.set_default_attn_processor()
504 |         assert all(type(proc) == AttnProcessor for proc in model.attn_processors.values())
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
505 |         with torch.no_grad():
506 |             if self.forward_requires_fresh_args:
    |

tests/models/test_modeling_common.py:512:20: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
511 |         model.set_attn_processor(AttnProcessor2_0())
512 |         assert all(type(proc) == AttnProcessor2_0 for proc in model.attn_processors.values())
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
513 |         with torch.no_grad():
514 |             if self.forward_requires_fresh_args:
    |

tests/models/test_modeling_common.py:520:20: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
519 |         model.set_attn_processor(AttnProcessor())
520 |         assert all(type(proc) == AttnProcessor for proc in model.attn_processors.values())
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
521 |         with torch.no_grad():
522 |             if self.forward_requires_fresh_args:
    |

tests/pipelines/test_pipelines_common.py:772:21: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
770 |             if hasattr(component, "attn_processors"):
771 |                 assert all(
772 |                     type(proc) == AttnProcessor for proc in component.attn_processors.values()
    |                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
773 |                 ), "`from_pipe` changed the attention processor in original pipeline."
    |

tests/schedulers/test_schedulers.py:827:16: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
825 |         scheduler_loaded = DDIMScheduler.from_pretrained(f"{USER}/{self.repo_id}")
826 |
827 |         assert type(scheduler) == type(scheduler_loaded)
    |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
828 |
829 |         # Reset repo
    |

tests/schedulers/test_schedulers.py:838:16: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
836 |         scheduler_loaded = DDIMScheduler.from_pretrained(f"{USER}/{self.repo_id}")
837 |
838 |         assert type(scheduler) == type(scheduler_loaded)
    |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
839 |
840 |         # Reset repo
    |

tests/schedulers/test_schedulers.py:854:16: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
852 |         scheduler_loaded = DDIMScheduler.from_pretrained(self.org_repo_id)
853 |
854 |         assert type(scheduler) == type(scheduler_loaded)
    |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
855 |
856 |         # Reset repo
    |

tests/schedulers/test_schedulers.py:865:16: E721 Use `is` and `is not` for type comparisons, or `isinstance()` for isinstance checks
    |
863 |         scheduler_loaded = DDIMScheduler.from_pretrained(self.org_repo_id)
864 |
865 |         assert type(scheduler) == type(scheduler_loaded)
    |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E721
866 |
867 |         # Reset repo

but the errors are not from this pr. I think this pr will not affect the make process.

[hlky]

Here's the error from the last run, link. The extra errors you're seeing are because of ruff version, we use 0.1.5, pip install ruff==0.1.5 or pip install -e ".[dev]" in Diffusers. Apologies for the inconvenience, your PR is from your main branch which seems to prevent us from pushing to it.

[zhangp365]

Here's the error from the last run, link. The extra errors you're seeing are because of ruff version, we use 0.1.5, pip install ruff==0.1.5 or pip install -e ".[dev]" in Diffusers. Apologies for the inconvenience, your PR is from your main branch which seems to prevent us from pushing to it.

Yes, after running pip install -e ".[dev]" in Diffusers, 'make style' is successful now. The error link was correct; at that time, there was a whitespace in an empty line, but I have since removed the empty line, so the issue no longer exists.

[hlky]

Thanks @zhangp365!