Spaces:
Runtime error
Runtime error
| # Copyright Philip Brown, ppbrown@github | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| ########################################################################### | |
| # This pipeline attempts to use a model that has SDXL vae, T5 text encoder, | |
| # and SDXL unet. | |
| # At the present time, there are no pretrained models that give pleasing | |
| # output. So as yet, (2025/06/10) this pipeline is somewhat of a tech | |
| # demo proving that the pieces can at least be put together. | |
| # Hopefully, it will encourage someone with the hardware available to | |
| # throw enough resources into training one up. | |
| from typing import Optional | |
| import torch.nn as nn | |
| from transformers import ( | |
| CLIPImageProcessor, | |
| CLIPTokenizer, | |
| CLIPVisionModelWithProjection, | |
| T5EncoderModel, | |
| ) | |
| from diffusers import DiffusionPipeline, StableDiffusionXLPipeline | |
| from diffusers.image_processor import VaeImageProcessor | |
| from diffusers.models import AutoencoderKL, UNet2DConditionModel | |
| from diffusers.schedulers import KarrasDiffusionSchedulers | |
| # Note: At this time, the intent is to use the T5 encoder mentioned | |
| # below, with zero changes. | |
| # Therefore, the model deliberately does not store the T5 encoder model bytes, | |
| # (Since they are not unique!) | |
| # but instead takes advantage of huggingface hub cache loading | |
| T5_NAME = "mcmonkey/google_t5-v1_1-xxl_encoderonly" | |
| # Caller is expected to load this, or equivalent, as model name for now | |
| # eg: pipe = StableDiffusionXL_T5Pipeline(SDXL_NAME) | |
| SDXL_NAME = "stabilityai/stable-diffusion-xl-base-1.0" | |
| class LinearWithDtype(nn.Linear): | |
| def dtype(self): | |
| return self.weight.dtype | |
| class StableDiffusionXL_T5Pipeline(StableDiffusionXLPipeline): | |
| _expected_modules = [ | |
| "vae", | |
| "unet", | |
| "scheduler", | |
| "tokenizer", | |
| "image_encoder", | |
| "feature_extractor", | |
| "t5_encoder", | |
| "t5_projection", | |
| "t5_pooled_projection", | |
| ] | |
| _optional_components = [ | |
| "image_encoder", | |
| "feature_extractor", | |
| "t5_encoder", | |
| "t5_projection", | |
| "t5_pooled_projection", | |
| ] | |
| def __init__( | |
| self, | |
| vae: AutoencoderKL, | |
| unet: UNet2DConditionModel, | |
| scheduler: KarrasDiffusionSchedulers, | |
| tokenizer: CLIPTokenizer, | |
| t5_encoder=None, | |
| t5_projection=None, | |
| t5_pooled_projection=None, | |
| image_encoder: CLIPVisionModelWithProjection = None, | |
| feature_extractor: CLIPImageProcessor = None, | |
| force_zeros_for_empty_prompt: bool = True, | |
| add_watermarker: Optional[bool] = None, | |
| ): | |
| DiffusionPipeline.__init__(self) | |
| if t5_encoder is None: | |
| self.t5_encoder = T5EncoderModel.from_pretrained(T5_NAME, torch_dtype=unet.dtype) | |
| else: | |
| self.t5_encoder = t5_encoder | |
| # ----- build T5 4096 => 2048 dim projection ----- | |
| if t5_projection is None: | |
| self.t5_projection = LinearWithDtype(4096, 2048) # trainable | |
| else: | |
| self.t5_projection = t5_projection | |
| self.t5_projection.to(dtype=unet.dtype) | |
| # ----- build T5 4096 => 1280 dim projection ----- | |
| if t5_pooled_projection is None: | |
| self.t5_pooled_projection = LinearWithDtype(4096, 1280) # trainable | |
| else: | |
| self.t5_pooled_projection = t5_pooled_projection | |
| self.t5_pooled_projection.to(dtype=unet.dtype) | |
| print("dtype of Linear is ", self.t5_projection.dtype) | |
| self.register_modules( | |
| vae=vae, | |
| unet=unet, | |
| scheduler=scheduler, | |
| tokenizer=tokenizer, | |
| t5_encoder=self.t5_encoder, | |
| t5_projection=self.t5_projection, | |
| t5_pooled_projection=self.t5_pooled_projection, | |
| image_encoder=image_encoder, | |
| feature_extractor=feature_extractor, | |
| ) | |
| self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) | |
| self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 | |
| self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) | |
| self.default_sample_size = ( | |
| self.unet.config.sample_size | |
| if hasattr(self, "unet") and self.unet is not None and hasattr(self.unet.config, "sample_size") | |
| else 128 | |
| ) | |
| self.watermark = None | |
| # Parts of original SDXL class complain if these attributes are not | |
| # at least PRESENT | |
| self.text_encoder = self.text_encoder_2 = None | |
| # ------------------------------------------------------------------ | |
| # Encode a text prompt (T5-XXL + 4096→2048 projection) | |
| # Returns exactly four tensors in the order SDXL’s __call__ expects. | |
| # ------------------------------------------------------------------ | |
| def encode_prompt( | |
| self, | |
| prompt, | |
| num_images_per_prompt: int = 1, | |
| do_classifier_free_guidance: bool = True, | |
| negative_prompt: str | None = None, | |
| **_, | |
| ): | |
| """ | |
| Returns | |
| ------- | |
| prompt_embeds : Tensor [B, T, 2048] | |
| negative_prompt_embeds : Tensor [B, T, 2048] | None | |
| pooled_prompt_embeds : Tensor [B, 1280] | |
| negative_pooled_prompt_embeds: Tensor [B, 1280] | None | |
| where B = batch * num_images_per_prompt | |
| """ | |
| # --- helper to tokenize on the pipeline’s device ---------------- | |
| def _tok(text: str): | |
| tok_out = self.tokenizer( | |
| text, | |
| return_tensors="pt", | |
| padding="max_length", | |
| max_length=self.tokenizer.model_max_length, | |
| truncation=True, | |
| ).to(self.device) | |
| return tok_out.input_ids, tok_out.attention_mask | |
| # ---------- positive stream ------------------------------------- | |
| ids, mask = _tok(prompt) | |
| h_pos = self.t5_encoder(ids, attention_mask=mask).last_hidden_state # [b, T, 4096] | |
| tok_pos = self.t5_projection(h_pos) # [b, T, 2048] | |
| pool_pos = self.t5_pooled_projection(h_pos.mean(dim=1)) # [b, 1280] | |
| # expand for multiple images per prompt | |
| tok_pos = tok_pos.repeat_interleave(num_images_per_prompt, 0) | |
| pool_pos = pool_pos.repeat_interleave(num_images_per_prompt, 0) | |
| # ---------- negative / CFG stream -------------------------------- | |
| if do_classifier_free_guidance: | |
| neg_text = "" if negative_prompt is None else negative_prompt | |
| ids_n, mask_n = _tok(neg_text) | |
| h_neg = self.t5_encoder(ids_n, attention_mask=mask_n).last_hidden_state | |
| tok_neg = self.t5_projection(h_neg) | |
| pool_neg = self.t5_pooled_projection(h_neg.mean(dim=1)) | |
| tok_neg = tok_neg.repeat_interleave(num_images_per_prompt, 0) | |
| pool_neg = pool_neg.repeat_interleave(num_images_per_prompt, 0) | |
| else: | |
| tok_neg = pool_neg = None | |
| # ----------------- final ordered return -------------------------- | |
| # 1) positive token embeddings | |
| # 2) negative token embeddings (or None) | |
| # 3) positive pooled embeddings | |
| # 4) negative pooled embeddings (or None) | |
| return tok_pos, tok_neg, pool_pos, pool_neg | |