recoilme commited on
Commit
ac79322
·
verified ·
1 Parent(s): 2eb35b4

Upload folder using huggingface_hub

Browse files
.ipynb_checkpoints/README-checkpoint.md ADDED
@@ -0,0 +1,116 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ pipeline_tag: text-to-image
4
+ ---
5
+ # Work / train in progress!
6
+ ![image](./promo.png)
7
+
8
+ ⚡️Waifu: efficient high-resolution waifu synthesis
9
+
10
+
11
+ waifu is a free text-to-image model that can efficiently generate images in 80 languages. Our goal is to create a small model without compromising on quality.
12
+
13
+ ## Core designs include:
14
+
15
+ (1) [**AuraDiffusion/16ch-vae**](https://huggingface.co/AuraDiffusion/16ch-vae): A fully open source 16ch VAE. Natively trained in fp16. \
16
+ (2) [**Linear DiT**](https://github.com/NVlabs/Sana): we use 1.6b DiT transformer with linear attention. \
17
+ (3) [**MEXMA-SigLIP**](https://huggingface.co/visheratin/mexma-siglip): MEXMA-SigLIP is a model that combines the [MEXMA](https://huggingface.co/facebook/MEXMA) multilingual text encoder and an image encoder from the [SigLIP](https://huggingface.co/timm/ViT-SO400M-14-SigLIP-384) model. This allows us to get a high-performance CLIP model for 80 languages.. \
18
+ (4) Other: we use Flow-Euler sampler, Adafactor-Fused optimizer and bf16 precision for training, and combine efficient caption labeling (MoonDream, CogVlM, Human, Gpt's) and danbooru tags to accelerate convergence.
19
+
20
+
21
+ ## Example
22
+
23
+ ```py
24
+ import torch
25
+ from diffusers import DiffusionPipeline
26
+
27
+ from transformers import XLMRobertaTokenizerFast,XLMRobertaModel
28
+ from diffusers import FlowMatchEulerDiscreteScheduler
29
+ from diffusers.models import AutoencoderKL
30
+ from diffusers import SanaTransformer2DModel
31
+
32
+ pipe_id = "AiArtLab/waifu-2b"
33
+ variant = "fp16"
34
+ # tokenizer
35
+ tokenizer = XLMRobertaTokenizerFast.from_pretrained(
36
+ pipe_id,
37
+ subfolder="tokenizer"
38
+ )
39
+
40
+ # text_encoder
41
+ text_encoder = XLMRobertaModel.from_pretrained(
42
+ pipe_id,
43
+ variant=variant,
44
+ subfolder="text_encoder",
45
+ add_pooling_layer=False
46
+ ).to("cuda")
47
+
48
+ # scheduler
49
+ scheduler = FlowMatchEulerDiscreteScheduler(shift=1.0)
50
+
51
+ # VAE
52
+ vae = AutoencoderKL.from_pretrained(
53
+ pipe_id,
54
+ variant=variant,
55
+ subfolder="vae"
56
+ ).to("cuda")
57
+
58
+ # Transformer
59
+ transformer = SanaTransformer2DModel.from_pretrained(
60
+ pipe_id,
61
+ variant=variant,
62
+ subfolder="transformer"
63
+ ).to("cuda")
64
+
65
+ # Pipeline
66
+ pipeline = DiffusionPipeline.from_pretrained(
67
+ pipe_id,
68
+ tokenizer=tokenizer,
69
+ text_encoder=text_encoder,
70
+ vae=vae,
71
+ transformer=transformer,
72
+ trust_remote_code=True,
73
+ ).to("cuda")
74
+ print(pipeline)
75
+
76
+ prompt = 'аниме девушка, waifu, يبتسم جنسيا , sur le fond de la tour Eiffel'
77
+ generator = torch.Generator(device="cuda").manual_seed(42)
78
+
79
+ image = pipeline(
80
+ prompt = prompt,
81
+ negative_prompt = "",
82
+ generator=generator,
83
+ )[0]
84
+
85
+ for img in image:
86
+ img.show()
87
+ img.save('waifu.png')
88
+
89
+ ```
90
+
91
+ ![image](./waifu.png)
92
+
93
+ ## Donations
94
+
95
+ We are a small GPU poor group of enthusiasts (current train budget ~$2k)
96
+
97
+ Please contact with us if you may provide some GPU's on training
98
+
99
+ DOGE: DEw2DR8C7BnF8GgcrfTzUjSnGkuMeJhg83
100
+
101
+ ![image](./1.png)
102
+ A fluffy domestic cat with piercing green eyes sits attentively in a sunlit room filled natural light streaming through large windows, its soft fur reflecting warm hues of orange from the golden glow casting across its sleek body and delicate features
103
+
104
+ ## Contacts
105
+
106
+ [recoilme](https://t.me/recoilme)
107
+
108
+ ## How to cite
109
+
110
+ ```bibtex
111
+ @misc{Waifu,
112
+ url = {[https://huggingface.co/AiArtLab/waifu-2b](https://huggingface.co/AiArtLab/waifu-2b)},
113
+ title = {waifu-2b},
114
+ author = {recoilme, muinez, femboysLover}
115
+ }
116
+ ```
.ipynb_checkpoints/model_index-checkpoint.json ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_class_name": ["pipeline_waifu", "WaifuPipeline"],
3
+ "_diffusers_version": "0.32.0.dev0",
4
+ "_name_or_path": "AiArtLab/waifu-2b",
5
+ "scheduler": [
6
+ "diffusers",
7
+ "FlowMatchEulerDiscreteScheduler"
8
+ ],
9
+ "text_encoder": [
10
+ "transformers",
11
+ "XLMRobertaModel"
12
+ ],
13
+ "tokenizer": [
14
+ "transformers",
15
+ "XLMRobertaTokenizerFast"
16
+ ],
17
+ "transformer": [
18
+ "diffusers",
19
+ "SanaTransformer2DModel"
20
+ ],
21
+ "vae": [
22
+ "diffusers",
23
+ "AutoencoderKL"
24
+ ]
25
+ }
.ipynb_checkpoints/pipeline_waifu-checkpoint.py CHANGED
@@ -1,23 +1,113 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import torch
 
 
2
  from diffusers import DiffusionPipeline
3
- from typing import Callable, Dict, List, Optional, Tuple, Union
4
-
5
- # waifu
6
- # tokenizer
7
- from transformers import XLMRobertaTokenizerFast
8
- # text_encoder
9
- from transformers import XLMRobertaModel
10
- # scheduler
11
- from diffusers import FlowMatchEulerDiscreteScheduler
12
- # VAE
13
- from diffusers.models import AutoencoderKL
14
- # Transformer
15
  from diffusers import SanaTransformer2DModel
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
16
 
17
 
18
  class WaifuPipeline(DiffusionPipeline):
19
  r"""
20
- Pipeline for text-to-image generation using [waifu](https://github.com/recoilme/waifu).
21
  """
22
 
23
  model_cpu_offload_seq = "text_encoder->transformer->vae"
@@ -37,9 +127,272 @@ class WaifuPipeline(DiffusionPipeline):
37
  tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
38
  )
39
 
40
- self.vae_scale_factor = 8
 
 
41
  self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor)
42
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43
  @torch.no_grad()
44
  def __call__(
45
  self,
@@ -60,11 +413,11 @@ class WaifuPipeline(DiffusionPipeline):
60
  negative_prompt_embeds: Optional[torch.Tensor] = None,
61
  negative_prompt_attention_mask: Optional[torch.Tensor] = None,
62
  output_type: Optional[str] = "pil",
63
- return_dict: bool = True,
64
  callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
65
  callback_on_step_end_tensor_inputs: List[str] = ["latents"],
66
  max_sequence_length: int = 512,
67
- ) -> Union[SanaPipelineOutput, Tuple]:
68
  """
69
  Function invoked when calling the pipeline for generation.
70
 
@@ -138,15 +491,15 @@ class WaifuPipeline(DiffusionPipeline):
138
  Examples:
139
 
140
  Returns:
141
- [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] or `tuple`:
142
- If `return_dict` is `True`, [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] is returned,
143
  otherwise a `tuple` is returned where the first element is a list with the generated images
144
  """
145
 
146
- if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
147
- callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
148
 
149
  # 1. Check inputs. Raise error if not correct
 
150
  self.check_inputs(
151
  prompt,
152
  height,
@@ -273,16 +626,16 @@ class WaifuPipeline(DiffusionPipeline):
273
  else:
274
  latents = latents.to(self.vae.dtype)
275
  image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
276
- if use_resolution_binning:
277
- image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height)
278
 
279
  if not output_type == "latent":
280
  image = self.image_processor.postprocess(image, output_type=output_type)
 
281
 
282
  # Offload all models
 
283
  self.maybe_free_model_hooks()
284
 
285
  if not return_dict:
286
  return (image,)
287
 
288
- return SanaPipelineOutput(images=image)
 
1
+ # Copyright 2024 PixArt-Sigma Authors and The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import inspect
16
+ from typing import Callable, Dict, List, Optional, Union
17
+
18
  import torch
19
+ from diffusers.image_processor import PixArtImageProcessor
20
+ from diffusers.utils.torch_utils import randn_tensor
21
  from diffusers import DiffusionPipeline
22
+ from transformers import XLMRobertaTokenizerFast,XLMRobertaModel
 
 
 
 
 
 
 
 
 
 
 
23
  from diffusers import SanaTransformer2DModel
24
+ from diffusers.models import AutoencoderKL
25
+ from diffusers import FlowMatchEulerDiscreteScheduler
26
+ from typing import List, Union
27
+ import numpy as np
28
+ import PIL.Image
29
+
30
+
31
+ EXAMPLE_DOC_STRING = """
32
+ Examples:
33
+ ```py
34
+ >>> import torch
35
+ >>> from diffusers import WaifuPipeline
36
+
37
+ >>> pipe = WaifuPipeline.from_pretrained(
38
+ ... "AiArtLab/waifu-2b"
39
+ ... )
40
+ >>> pipe.to("cuda")
41
+
42
+ >>> image = pipe(prompt='a cyberpunk cat with a neon sign that says "Sana"')[0]
43
+ >>> image[0].save("output.png")
44
+ ```
45
+ """
46
+
47
+
48
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
49
+ def retrieve_timesteps(
50
+ scheduler,
51
+ num_inference_steps: Optional[int] = None,
52
+ device: Optional[Union[str, torch.device]] = None,
53
+ timesteps: Optional[List[int]] = None,
54
+ sigmas: Optional[List[float]] = None,
55
+ **kwargs,
56
+ ):
57
+ r"""
58
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
59
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
60
+
61
+ Args:
62
+ scheduler (`SchedulerMixin`):
63
+ The scheduler to get timesteps from.
64
+ num_inference_steps (`int`):
65
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
66
+ must be `None`.
67
+ device (`str` or `torch.device`, *optional*):
68
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
69
+ timesteps (`List[int]`, *optional*):
70
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
71
+ `num_inference_steps` and `sigmas` must be `None`.
72
+ sigmas (`List[float]`, *optional*):
73
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
74
+ `num_inference_steps` and `timesteps` must be `None`.
75
+
76
+ Returns:
77
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
78
+ second element is the number of inference steps.
79
+ """
80
+ if timesteps is not None and sigmas is not None:
81
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
82
+ if timesteps is not None:
83
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
84
+ if not accepts_timesteps:
85
+ raise ValueError(
86
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
87
+ f" timestep schedules. Please check whether you are using the correct scheduler."
88
+ )
89
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
90
+ timesteps = scheduler.timesteps
91
+ num_inference_steps = len(timesteps)
92
+ elif sigmas is not None:
93
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
94
+ if not accept_sigmas:
95
+ raise ValueError(
96
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
97
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
98
+ )
99
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
100
+ timesteps = scheduler.timesteps
101
+ num_inference_steps = len(timesteps)
102
+ else:
103
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
104
+ timesteps = scheduler.timesteps
105
+ return timesteps, num_inference_steps
106
 
107
 
108
  class WaifuPipeline(DiffusionPipeline):
109
  r"""
110
+ Pipeline for text-to-image generation using [Sana](https://huggingface.co/papers/2410.10629).
111
  """
112
 
113
  model_cpu_offload_seq = "text_encoder->transformer->vae"
 
127
  tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
128
  )
129
 
130
+ self.vae_scale_factor = (
131
+ 8
132
+ )
133
  self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor)
134
 
135
+ def encode_prompt(
136
+ self,
137
+ prompt: Union[str, List[str]],
138
+ do_classifier_free_guidance: bool = True,
139
+ negative_prompt: str = "",
140
+ num_images_per_prompt: int = 1,
141
+ device: Optional[torch.device] = None,
142
+ prompt_embeds: Optional[torch.Tensor] = None,
143
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
144
+ prompt_attention_mask: Optional[torch.Tensor] = None,
145
+ negative_prompt_attention_mask: Optional[torch.Tensor] = None,
146
+ max_sequence_length: int = 512,
147
+ ):
148
+ r"""
149
+ Encodes the prompt into text encoder hidden states.
150
+
151
+ Args:
152
+ prompt (`str` or `List[str]`, *optional*):
153
+ prompt to be encoded
154
+ negative_prompt (`str` or `List[str]`, *optional*):
155
+ The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
156
+ instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
157
+ PixArt-Alpha, this should be "".
158
+ do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
159
+ whether to use classifier free guidance or not
160
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
161
+ number of images that should be generated per prompt
162
+ device: (`torch.device`, *optional*):
163
+ torch device to place the resulting embeddings on
164
+ prompt_embeds (`torch.Tensor`, *optional*):
165
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
166
+ provided, text embeddings will be generated from `prompt` input argument.
167
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
168
+ Pre-generated negative text embeddings. For Sana, it's should be the embeddings of the "" string.
169
+ max_sequence_length (`int`, defaults to 512): Maximum sequence length to use for the prompt.
170
+ """
171
+
172
+ if device is None:
173
+ device = self._execution_device
174
+
175
+ if prompt is not None and isinstance(prompt, str):
176
+ batch_size = 1
177
+ elif prompt is not None and isinstance(prompt, list):
178
+ batch_size = len(prompt)
179
+ else:
180
+ batch_size = prompt_embeds.shape[0]
181
+
182
+ if self.tokenizer is not None:
183
+ self.tokenizer.padding_side = "right"
184
+
185
+ max_length = max_sequence_length
186
+ select_index = [0] + list(range(-max_length + 1, 0))
187
+
188
+ if prompt_embeds is None:
189
+ prompt = self._text_preprocessing(prompt)
190
+
191
+ max_length_all = max_length
192
+
193
+ text_inputs = self.tokenizer(
194
+ prompt,
195
+ padding="max_length",
196
+ max_length=max_length_all,
197
+ truncation=True,
198
+ add_special_tokens=True,
199
+ return_tensors="pt",
200
+ )
201
+ text_input_ids = text_inputs.input_ids
202
+
203
+ prompt_attention_mask = text_inputs.attention_mask
204
+ prompt_attention_mask = prompt_attention_mask.to(device)
205
+
206
+ prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)
207
+ prompt_embeds = prompt_embeds[0][:, select_index]
208
+ prompt_attention_mask = prompt_attention_mask[:, select_index]
209
+
210
+ if self.transformer is not None:
211
+ dtype = self.transformer.dtype
212
+ elif self.text_encoder is not None:
213
+ dtype = self.text_encoder.dtype
214
+ else:
215
+ dtype = None
216
+
217
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
218
+
219
+ bs_embed, seq_len, _ = prompt_embeds.shape
220
+ # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
221
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
222
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
223
+ prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
224
+ prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
225
+
226
+ # get unconditional embeddings for classifier free guidance
227
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
228
+ #print("do_classifier_free_guidance and negative_prompt_embeds is None")
229
+ uncond_tokens = [negative_prompt] * batch_size if isinstance(negative_prompt, str) else negative_prompt
230
+ uncond_tokens = self._text_preprocessing(uncond_tokens)
231
+ max_length = prompt_embeds.shape[1]
232
+ uncond_input = self.tokenizer(
233
+ uncond_tokens,
234
+ padding="max_length",
235
+ max_length=max_length,
236
+ truncation=True,
237
+ return_attention_mask=True,
238
+ add_special_tokens=True,
239
+ return_tensors="pt",
240
+ )
241
+ negative_prompt_attention_mask = uncond_input.attention_mask
242
+ negative_prompt_attention_mask = negative_prompt_attention_mask.to(device)
243
+
244
+ negative_prompt_embeds = self.text_encoder(
245
+ uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask
246
+ )
247
+ negative_prompt_embeds = negative_prompt_embeds[0]
248
+
249
+ if do_classifier_free_guidance:
250
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
251
+ seq_len = negative_prompt_embeds.shape[1]
252
+
253
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
254
+
255
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
256
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
257
+
258
+ negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1)
259
+ negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
260
+ else:
261
+ negative_prompt_embeds = None
262
+ negative_prompt_attention_mask = None
263
+
264
+ return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
265
+
266
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
267
+ def prepare_extra_step_kwargs(self, generator, eta):
268
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
269
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
270
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
271
+ # and should be between [0, 1]
272
+
273
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
274
+ extra_step_kwargs = {}
275
+ if accepts_eta:
276
+ extra_step_kwargs["eta"] = eta
277
+
278
+ # check if the scheduler accepts generator
279
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
280
+ if accepts_generator:
281
+ extra_step_kwargs["generator"] = generator
282
+ return extra_step_kwargs
283
+
284
+ def check_inputs(
285
+ self,
286
+ prompt,
287
+ height,
288
+ width,
289
+ callback_on_step_end_tensor_inputs=None,
290
+ negative_prompt=None,
291
+ prompt_embeds=None,
292
+ negative_prompt_embeds=None,
293
+ prompt_attention_mask=None,
294
+ negative_prompt_attention_mask=None,
295
+ ):
296
+ if height % 64 != 0 or width % 64 != 0:
297
+ raise ValueError(f"`height` and `width` have to be divisible by 64 but are {height} and {width}.")
298
+
299
+ if callback_on_step_end_tensor_inputs is not None and not all(
300
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
301
+ ):
302
+ raise ValueError(
303
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
304
+ )
305
+
306
+ if prompt is not None and prompt_embeds is not None:
307
+ raise ValueError(
308
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
309
+ " only forward one of the two."
310
+ )
311
+ elif prompt is None and prompt_embeds is None:
312
+ raise ValueError(
313
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
314
+ )
315
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
316
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
317
+
318
+ if prompt is not None and negative_prompt_embeds is not None:
319
+ raise ValueError(
320
+ f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
321
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
322
+ )
323
+
324
+ if negative_prompt is not None and negative_prompt_embeds is not None:
325
+ raise ValueError(
326
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
327
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
328
+ )
329
+
330
+ if prompt_embeds is not None and prompt_attention_mask is None:
331
+ raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
332
+
333
+ if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
334
+ raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
335
+
336
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
337
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
338
+ raise ValueError(
339
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
340
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
341
+ f" {negative_prompt_embeds.shape}."
342
+ )
343
+ if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
344
+ raise ValueError(
345
+ "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
346
+ f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
347
+ f" {negative_prompt_attention_mask.shape}."
348
+ )
349
+
350
+ def _text_preprocessing(self, text):
351
+
352
+ if not isinstance(text, (tuple, list)):
353
+ text = [text]
354
+
355
+ def process(text: str):
356
+ text = text.lower().strip()
357
+ return text
358
+
359
+ return [process(t) for t in text]
360
+
361
+ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
362
+ if latents is not None:
363
+ return latents.to(device=device, dtype=dtype)
364
+
365
+ shape = (
366
+ batch_size,
367
+ num_channels_latents,
368
+ int(height) // self.vae_scale_factor,
369
+ int(width) // self.vae_scale_factor,
370
+ )
371
+ if isinstance(generator, list) and len(generator) != batch_size:
372
+ raise ValueError(
373
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
374
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
375
+ )
376
+
377
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
378
+ return latents
379
+
380
+ @property
381
+ def guidance_scale(self):
382
+ return self._guidance_scale
383
+
384
+ @property
385
+ def do_classifier_free_guidance(self):
386
+ return self._guidance_scale > 1.0
387
+
388
+ @property
389
+ def num_timesteps(self):
390
+ return self._num_timesteps
391
+
392
+ @property
393
+ def interrupt(self):
394
+ return self._interrupt
395
+
396
  @torch.no_grad()
397
  def __call__(
398
  self,
 
413
  negative_prompt_embeds: Optional[torch.Tensor] = None,
414
  negative_prompt_attention_mask: Optional[torch.Tensor] = None,
415
  output_type: Optional[str] = "pil",
416
+ return_dict: bool = False,
417
  callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
418
  callback_on_step_end_tensor_inputs: List[str] = ["latents"],
419
  max_sequence_length: int = 512,
420
+ ) -> Union[List[PIL.Image.Image], np.ndarray]:
421
  """
422
  Function invoked when calling the pipeline for generation.
423
 
 
491
  Examples:
492
 
493
  Returns:
494
+ Union[List[PIL.Image.Image], np.ndarray] is returned,
 
495
  otherwise a `tuple` is returned where the first element is a list with the generated images
496
  """
497
 
498
+ # if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
499
+ # callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
500
 
501
  # 1. Check inputs. Raise error if not correct
502
+
503
  self.check_inputs(
504
  prompt,
505
  height,
 
626
  else:
627
  latents = latents.to(self.vae.dtype)
628
  image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
 
 
629
 
630
  if not output_type == "latent":
631
  image = self.image_processor.postprocess(image, output_type=output_type)
632
+ #image = numpy_to_pil(image)
633
 
634
  # Offload all models
635
+ #print("Offload all models 4")
636
  self.maybe_free_model_hooks()
637
 
638
  if not return_dict:
639
  return (image,)
640
 
641
+ return Union[List[PIL.Image.Image], np.ndarray]
.ipynb_checkpoints/test-checkpoint.ipynb ADDED
The diff for this file is too large to render. See raw diff
 
.ipynb_checkpoints/waifu-checkpoint.png ADDED
1.png CHANGED
README.md CHANGED
@@ -2,7 +2,8 @@
2
  license: apache-2.0
3
  pipeline_tag: text-to-image
4
  ---
5
- # Work / train in progress
 
6
 
7
  ⚡️Waifu: efficient high-resolution waifu synthesis
8
 
 
2
  license: apache-2.0
3
  pipeline_tag: text-to-image
4
  ---
5
+ # Work / train in progress!
6
+ ![image](./promo.png)
7
 
8
  ⚡️Waifu: efficient high-resolution waifu synthesis
9
 
Untitled.ipynb CHANGED
@@ -172,7 +172,7 @@
172
  "name": "python",
173
  "nbconvert_exporter": "python",
174
  "pygments_lexer": "ipython3",
175
- "version": "3.11.6"
176
  }
177
  },
178
  "nbformat": 4,
 
172
  "name": "python",
173
  "nbconvert_exporter": "python",
174
  "pygments_lexer": "ipython3",
175
+ "version": "3.10.12"
176
  }
177
  },
178
  "nbformat": 4,
__pycache__/pipeline_waifu.cpython-310.pyc ADDED
Binary file (20.3 kB). View file
 
pipeline_waifu.py CHANGED
@@ -115,11 +115,11 @@ class WaifuPipeline(DiffusionPipeline):
115
 
116
  def __init__(
117
  self,
118
- tokenizer: None,
119
- text_encoder: None,
120
- vae: None,
121
- transformer: None,
122
- scheduler: None,
123
  ):
124
  super().__init__()
125
 
 
115
 
116
  def __init__(
117
  self,
118
+ tokenizer: XLMRobertaTokenizerFast,
119
+ text_encoder: XLMRobertaModel,
120
+ vae: AutoencoderKL,
121
+ transformer: SanaTransformer2DModel,
122
+ scheduler: FlowMatchEulerDiscreteScheduler,
123
  ):
124
  super().__init__()
125
 
promo.png CHANGED

Git LFS Details

  • SHA256: 1505648f1fd5b4f352bb621d12942db147f56ac9fce42069ce5e5e007b7a1d98
  • Pointer size: 132 Bytes
  • Size of remote file: 1.81 MB

Git LFS Details

  • SHA256: 12a5ea2c007c532b68f6ce3f60572ce05e3e08edc0f814b077822937e727e500
  • Pointer size: 132 Bytes
  • Size of remote file: 1.69 MB
test.ipynb CHANGED
The diff for this file is too large to render. See raw diff
 
transformer/diffusion_pytorch_model.fp16.safetensors CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:7b8de92385645ac8c8e4a6ac9072ecc832a9f639d6acec5726ab87943f880791
3
  size 3203093344
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2d41fc249d188765f985a9502e806a8ea239580ba95387ed323dfb6304a48515
3
  size 3203093344
waifu.png CHANGED