Fix generator determinism: forward generator through scheduler steps and seeded noise
Browse files- ADM-G-256/pipeline.py +10 -6
- ADM-G-512/pipeline.py +13 -9
ADM-G-256/pipeline.py
CHANGED
|
@@ -1,3 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
#
|
| 3 |
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
@@ -19,7 +25,6 @@ import torch
|
|
| 19 |
|
| 20 |
from diffusers.image_processor import VaeImageProcessor
|
| 21 |
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
|
| 22 |
-
from diffusers.schedulers import KarrasDiffusionSchedulers
|
| 23 |
from diffusers.utils import replace_example_docstring
|
| 24 |
from diffusers.utils.torch_utils import randn_tensor
|
| 25 |
|
|
@@ -44,7 +49,6 @@ EXAMPLE_DOC_STRING = """
|
|
| 44 |
```
|
| 45 |
"""
|
| 46 |
|
| 47 |
-
|
| 48 |
class ADMPipeline(DiffusionPipeline):
|
| 49 |
r"""ADM/ADM-G pipeline compatible with Diffusers custom pipeline loading."""
|
| 50 |
|
|
@@ -54,7 +58,7 @@ class ADMPipeline(DiffusionPipeline):
|
|
| 54 |
def __init__(
|
| 55 |
self,
|
| 56 |
unet,
|
| 57 |
-
scheduler
|
| 58 |
classifier: Optional[Any] = None,
|
| 59 |
id2label: Optional[Dict[str, str]] = None,
|
| 60 |
null_class_id: int = 1000,
|
|
@@ -92,7 +96,7 @@ class ADMPipeline(DiffusionPipeline):
|
|
| 92 |
|
| 93 |
@staticmethod
|
| 94 |
def prepare_extra_step_kwargs(
|
| 95 |
-
scheduler
|
| 96 |
generator: Optional[Union[torch.Generator, List[torch.Generator]]],
|
| 97 |
eta: float,
|
| 98 |
) -> Dict[str, Any]:
|
|
@@ -112,7 +116,7 @@ class ADMPipeline(DiffusionPipeline):
|
|
| 112 |
def _prepare_model_output_for_scheduler(
|
| 113 |
model_output: torch.Tensor,
|
| 114 |
channels: int,
|
| 115 |
-
scheduler
|
| 116 |
) -> torch.Tensor:
|
| 117 |
if model_output.shape[1] != 2 * channels:
|
| 118 |
return model_output
|
|
@@ -266,4 +270,4 @@ class ADMPipeline(DiffusionPipeline):
|
|
| 266 |
self.maybe_free_model_hooks()
|
| 267 |
if not return_dict:
|
| 268 |
return (image,)
|
| 269 |
-
return ImagePipelineOutput(images=image)
|
|
|
|
| 1 |
+
"""Hub custom pipeline: ADMPipeline.
|
| 2 |
+
Load with native Hugging Face diffusers and trust_remote_code=True.
|
| 3 |
+
"""
|
| 4 |
+
|
| 5 |
+
from __future__ import annotations
|
| 6 |
+
|
| 7 |
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 8 |
#
|
| 9 |
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
|
| 25 |
|
| 26 |
from diffusers.image_processor import VaeImageProcessor
|
| 27 |
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
|
|
|
|
| 28 |
from diffusers.utils import replace_example_docstring
|
| 29 |
from diffusers.utils.torch_utils import randn_tensor
|
| 30 |
|
|
|
|
| 49 |
```
|
| 50 |
"""
|
| 51 |
|
|
|
|
| 52 |
class ADMPipeline(DiffusionPipeline):
|
| 53 |
r"""ADM/ADM-G pipeline compatible with Diffusers custom pipeline loading."""
|
| 54 |
|
|
|
|
| 58 |
def __init__(
|
| 59 |
self,
|
| 60 |
unet,
|
| 61 |
+
scheduler,
|
| 62 |
classifier: Optional[Any] = None,
|
| 63 |
id2label: Optional[Dict[str, str]] = None,
|
| 64 |
null_class_id: int = 1000,
|
|
|
|
| 96 |
|
| 97 |
@staticmethod
|
| 98 |
def prepare_extra_step_kwargs(
|
| 99 |
+
scheduler,
|
| 100 |
generator: Optional[Union[torch.Generator, List[torch.Generator]]],
|
| 101 |
eta: float,
|
| 102 |
) -> Dict[str, Any]:
|
|
|
|
| 116 |
def _prepare_model_output_for_scheduler(
|
| 117 |
model_output: torch.Tensor,
|
| 118 |
channels: int,
|
| 119 |
+
scheduler,
|
| 120 |
) -> torch.Tensor:
|
| 121 |
if model_output.shape[1] != 2 * channels:
|
| 122 |
return model_output
|
|
|
|
| 270 |
self.maybe_free_model_hooks()
|
| 271 |
if not return_dict:
|
| 272 |
return (image,)
|
| 273 |
+
return ImagePipelineOutput(images=image)
|
ADM-G-512/pipeline.py
CHANGED
|
@@ -1,3 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 2 |
#
|
| 3 |
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
@@ -19,7 +25,6 @@ import torch
|
|
| 19 |
|
| 20 |
from diffusers.image_processor import VaeImageProcessor
|
| 21 |
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
|
| 22 |
-
from diffusers.schedulers import KarrasDiffusionSchedulers
|
| 23 |
from diffusers.utils import replace_example_docstring
|
| 24 |
from diffusers.utils.torch_utils import randn_tensor
|
| 25 |
|
|
@@ -30,7 +35,7 @@ EXAMPLE_DOC_STRING = """
|
|
| 30 |
>>> import torch
|
| 31 |
>>> from diffusers import DiffusionPipeline
|
| 32 |
|
| 33 |
-
>>> model_dir = Path("path/to/BiliSakura/ADM-diffusers/ADM-G-
|
| 34 |
>>> pipe = DiffusionPipeline.from_pretrained(
|
| 35 |
... str(model_dir),
|
| 36 |
... local_files_only=True,
|
|
@@ -40,11 +45,10 @@ EXAMPLE_DOC_STRING = """
|
|
| 40 |
... )
|
| 41 |
>>> pipe = pipe.to("cuda")
|
| 42 |
>>> class_id = pipe.get_label_ids("golden retriever")[0]
|
| 43 |
-
>>> image = pipe(class_labels=class_id, guidance_scale=
|
| 44 |
```
|
| 45 |
"""
|
| 46 |
|
| 47 |
-
|
| 48 |
class ADMPipeline(DiffusionPipeline):
|
| 49 |
r"""ADM/ADM-G pipeline compatible with Diffusers custom pipeline loading."""
|
| 50 |
|
|
@@ -54,7 +58,7 @@ class ADMPipeline(DiffusionPipeline):
|
|
| 54 |
def __init__(
|
| 55 |
self,
|
| 56 |
unet,
|
| 57 |
-
scheduler
|
| 58 |
classifier: Optional[Any] = None,
|
| 59 |
id2label: Optional[Dict[str, str]] = None,
|
| 60 |
null_class_id: int = 1000,
|
|
@@ -92,7 +96,7 @@ class ADMPipeline(DiffusionPipeline):
|
|
| 92 |
|
| 93 |
@staticmethod
|
| 94 |
def prepare_extra_step_kwargs(
|
| 95 |
-
scheduler
|
| 96 |
generator: Optional[Union[torch.Generator, List[torch.Generator]]],
|
| 97 |
eta: float,
|
| 98 |
) -> Dict[str, Any]:
|
|
@@ -112,7 +116,7 @@ class ADMPipeline(DiffusionPipeline):
|
|
| 112 |
def _prepare_model_output_for_scheduler(
|
| 113 |
model_output: torch.Tensor,
|
| 114 |
channels: int,
|
| 115 |
-
scheduler
|
| 116 |
) -> torch.Tensor:
|
| 117 |
if model_output.shape[1] != 2 * channels:
|
| 118 |
return model_output
|
|
@@ -139,7 +143,7 @@ class ADMPipeline(DiffusionPipeline):
|
|
| 139 |
height: Optional[int] = None,
|
| 140 |
width: Optional[int] = None,
|
| 141 |
num_inference_steps: int = 250,
|
| 142 |
-
guidance_scale: float =
|
| 143 |
classifier_guidance_scale: float = 0.0,
|
| 144 |
eta: float = 0.0,
|
| 145 |
clip_denoised: bool = True,
|
|
@@ -266,4 +270,4 @@ class ADMPipeline(DiffusionPipeline):
|
|
| 266 |
self.maybe_free_model_hooks()
|
| 267 |
if not return_dict:
|
| 268 |
return (image,)
|
| 269 |
-
return ImagePipelineOutput(images=image)
|
|
|
|
| 1 |
+
"""Hub custom pipeline: ADMPipeline.
|
| 2 |
+
Load with native Hugging Face diffusers and trust_remote_code=True.
|
| 3 |
+
"""
|
| 4 |
+
|
| 5 |
+
from __future__ import annotations
|
| 6 |
+
|
| 7 |
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 8 |
#
|
| 9 |
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
|
| 25 |
|
| 26 |
from diffusers.image_processor import VaeImageProcessor
|
| 27 |
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
|
|
|
|
| 28 |
from diffusers.utils import replace_example_docstring
|
| 29 |
from diffusers.utils.torch_utils import randn_tensor
|
| 30 |
|
|
|
|
| 35 |
>>> import torch
|
| 36 |
>>> from diffusers import DiffusionPipeline
|
| 37 |
|
| 38 |
+
>>> model_dir = Path("path/to/BiliSakura/ADM-diffusers/ADM-G-256")
|
| 39 |
>>> pipe = DiffusionPipeline.from_pretrained(
|
| 40 |
... str(model_dir),
|
| 41 |
... local_files_only=True,
|
|
|
|
| 45 |
... )
|
| 46 |
>>> pipe = pipe.to("cuda")
|
| 47 |
>>> class_id = pipe.get_label_ids("golden retriever")[0]
|
| 48 |
+
>>> image = pipe(class_labels=class_id, guidance_scale=1.0).images[0]
|
| 49 |
```
|
| 50 |
"""
|
| 51 |
|
|
|
|
| 52 |
class ADMPipeline(DiffusionPipeline):
|
| 53 |
r"""ADM/ADM-G pipeline compatible with Diffusers custom pipeline loading."""
|
| 54 |
|
|
|
|
| 58 |
def __init__(
|
| 59 |
self,
|
| 60 |
unet,
|
| 61 |
+
scheduler,
|
| 62 |
classifier: Optional[Any] = None,
|
| 63 |
id2label: Optional[Dict[str, str]] = None,
|
| 64 |
null_class_id: int = 1000,
|
|
|
|
| 96 |
|
| 97 |
@staticmethod
|
| 98 |
def prepare_extra_step_kwargs(
|
| 99 |
+
scheduler,
|
| 100 |
generator: Optional[Union[torch.Generator, List[torch.Generator]]],
|
| 101 |
eta: float,
|
| 102 |
) -> Dict[str, Any]:
|
|
|
|
| 116 |
def _prepare_model_output_for_scheduler(
|
| 117 |
model_output: torch.Tensor,
|
| 118 |
channels: int,
|
| 119 |
+
scheduler,
|
| 120 |
) -> torch.Tensor:
|
| 121 |
if model_output.shape[1] != 2 * channels:
|
| 122 |
return model_output
|
|
|
|
| 143 |
height: Optional[int] = None,
|
| 144 |
width: Optional[int] = None,
|
| 145 |
num_inference_steps: int = 250,
|
| 146 |
+
guidance_scale: float = 1.0,
|
| 147 |
classifier_guidance_scale: float = 0.0,
|
| 148 |
eta: float = 0.0,
|
| 149 |
clip_denoised: bool = True,
|
|
|
|
| 270 |
self.maybe_free_model_hooks()
|
| 271 |
if not return_dict:
|
| 272 |
return (image,)
|
| 273 |
+
return ImagePipelineOutput(images=image)
|