Merge pull request #5 from asomoza/main

LCM and SSD-1B support added

CHANGELOG.MD

@@ -0,0 +1,13 @@
+# Changelog
+
+## 0.1.0 - 2023-11-09
+
+- initial release
+- Windows installer
+
+## 0.1.1 - 2023-11-10
+
+- Support for LCM models
+- Support for SSD-1B models
+- Support for LMC LoRAs
+- Diffusers 0.23 dependency required

README.md

@@ -38,11 +38,13 @@ It is highly recommended to use the included VAE with the FP16 fix, since the VA
 - Will completely run offline after the first installation.
 - Powerfull features only avalaible as a desktop application.
 - Easy sharing of models and Lora's metadata since the information its stored in each model, including sample image, sample generation, triggers and tags for filtering.
+- Latent Consistency Models (LCM) and LoRAs for fast inference.
+- Segmind Stable Diffusion (SSD-1B) models for VRAM savings.
 
 ## Limitations
 
 - Only runs with Stable Diffusion XL models.
-- It has the default 75 CLIP token limitation for the prompts.
+- It has the default CLIP 75 token limitation for the prompts.
 
 You can read why [here](https://github.com/ZCode-opensource/image-artisan-xl/blob/main/EXPLANATIONS.MD).
 

pyproject.toml

@@ -4,15 +4,15 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "iartisanxl"
-version = "0.0.1"
+version = "0.1.1"
 authors = [
  { name="Alvaro Somoza", email="[email protected]" },
 ]
 description = "Dekstop application for generating images using Stable Diffusion."
 requires-python = ">=3.11"
 dependencies = [
  "accelerate>=0.24.1 ",
- "diffusers>=0.22.0",
+ "diffusers>=0.23.0",
  "Pillow>=9.3.0",
  "PyOpenGL",
  "PyOpenGL_accelerate",

src/iartisanxl/convert_model/convert_functions.py

@@ -763,6 +763,10 @@ def convert_ldm_unet_checkpoint(
  "label_emb.0.2.bias"
  ]
 
+ # Relevant to StableDiffusionUpscalePipeline
+ if "num_class_embeds" in config:
+ new_checkpoint["class_embedding.weight"] = unet_state_dict["label_emb.weight"]
+
  new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"]
  new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"]
 
@@ -845,6 +849,7 @@ def convert_ldm_unet_checkpoint(
 
  if len(attentions):
  paths = renew_attention_paths(attentions)
+
  meta_path = {
  "old": f"input_blocks.{i}.1",
  "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}",

src/iartisanxl/generation/schedulers/schedulers.py

@@ -15,6 +15,7 @@
  DEISMultistepScheduler,
  DDPMScheduler,
  DPMSolverSDEScheduler,
+ LCMScheduler,
 )
 
 
@@ -68,4 +69,5 @@ class Scheduler:
  Scheduler("LMS Karras", LMSDiscreteScheduler, dict(use_karras_sigmas=True)),
  Scheduler("Euler Ancestral", EulerAncestralDiscreteScheduler, dict()),
  Scheduler("KDPM 2 Ancestral", KDPM2AncestralDiscreteScheduler, dict()),
+ Scheduler("LCM", LCMScheduler, dict()),
 ]

src/iartisanxl/pipelines/txt_pipeline.py

@@ -358,6 +358,20 @@ def _get_add_time_ids(
  add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
  return add_time_ids
 
+ def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32):
+ assert len(w.shape) == 1
+ w = w * 1000.0
+
+ half_dim = embedding_dim // 2
+ emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
+ emb = w.to(dtype)[:, None] * emb[None, :]
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+ if embedding_dim % 2 == 1: # zero pad
+ emb = torch.nn.functional.pad(emb, (0, 1)) # pylint: disable=not-callable
+ assert emb.shape == (w.shape[0], embedding_dim)
+ return emb
+
  @torch.no_grad()
  def __call__(
  self,
@@ -387,6 +401,10 @@ def __call__(
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  self.logger.debug("Using device: %s", device)
 
+ do_classifier_free_guidance = False
+ if guidance_scale > 1:
+ do_classifier_free_guidance = True
+
  status_update("Encoding the prompt...")
  text_encoder_lora_scale = (
  cross_attention_kwargs.get("scale", None)
@@ -455,16 +473,25 @@ def __call__(
  negative_add_time_ids = add_time_ids
 
  status_update("Preparing emdeddings...")
- prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
- add_text_embeds = torch.cat(
- [negative_pooled_prompt_embeds, add_text_embeds], dim=0
- )
- add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+ if do_classifier_free_guidance:
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+ add_text_embeds = torch.cat(
+ [negative_pooled_prompt_embeds, add_text_embeds], dim=0
+ )
+ add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
 
  prompt_embeds = prompt_embeds.to(device)
  add_text_embeds = add_text_embeds.to(device)
  add_time_ids = add_time_ids.to(device)
 
+ timestep_cond = None
+ if self.unet.config.time_cond_proj_dim is not None:
+ guidance_scale_tensor = torch.tensor(guidance_scale - 1).repeat(1)
+ timestep_cond = self.get_guidance_scale_embedding(
+ guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+ ).to(device=device, dtype=latents.dtype)
+
  status_update("Generating image...")
  num_warmup_steps = max(
  len(timesteps) - num_inference_steps * self.scheduler.order, 0
@@ -478,7 +505,9 @@ def __call__(
  return
 
  # expand the latents
- latent_model_input = torch.cat([latents] * 2)
+ latent_model_input = (
+ torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+ )
  latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
 
  # predict the noise residual
@@ -491,6 +520,7 @@ def __call__(
  latent_model_input,
  t,
  encoder_hidden_states=prompt_embeds,
+ timestep_cond=timestep_cond,
  cross_attention_kwargs=cross_attention_kwargs,
  added_cond_kwargs=added_cond_kwargs,
  return_dict=False,
@@ -501,10 +531,11 @@ def __call__(
  return
 
  # perform guidance
- noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
- noise_pred = noise_pred_uncond + guidance_scale * (
- noise_pred_text - noise_pred_uncond
- )
+ if do_classifier_free_guidance:
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+ noise_pred = noise_pred_uncond + guidance_scale * (
+ noise_pred_text - noise_pred_uncond
+ )
 
  # compute the previous noisy sample x_t -> x_t-1
  latents = self.scheduler.step(