fixes

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 
-from diffusers import AutoPipelineForInpainting
+from diffusers import AutoPipelineForInpainting, StableDiffusionXLPipeline
 import diffusers
 from share_btn import community_icon_html, loading_icon_html, share_js
 from sdxl import gen_sdxl_simplified_interface
@@ -10,9 +10,14 @@ from sdxl_models import SDXLUNet, SDXLVae, SDXLControlNetPreEncodedControlnetCon
 device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to(device)
 
-comparing_unet = SDXLUNet.load_fp16(device=device)
-comparing_vae = SDXLVae.load_fp16_fix(device=device)
-comparing_controlnet = SDXLControlNetPreEncodedControlnetCond.load("", device="cuda") # TODO - upload checkpoint
+# TODO - just download individual files
+# StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", variant="fp16") # download weights
+comparing_unet = SDXLUNet.load("/admin/home/william/.cache/huggingface/hub/models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/76d28af79639c28a79fa5c6c6468febd3490a37e/unet/diffusion_pytorch_model.fp16.safetensors", device=device)
+# comparing_vae = SDXLVae.load("/admin/home/william/.cache/huggingface/hub/models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/76d28af79639c28a79fa5c6c6468febd3490a37e/vae/diffusion_pytorch_model.fp16.safetensors", device=device)
+comparing_vae = SDXLVae.load("/admin/home/william/.cache/huggingface/hub/models--madebyollin--sdxl-vae-fp16-fix/snapshots/4df413ca49271c25289a6482ab97a433f8117d15/diffusion_pytorch_model.safetensors", device=device)
+comparing_vae.to(torch.float16)
+# comparing_controlnet = SDXLControlNetPreEncodedControlnetCond.load("/fsx/william/diffusers-utils/output/sdxl_controlnet_inpaint_pre_encoded_controlnet_cond/checkpoint-200000/controlnet/diffusion_pytorch_model.safetensors", device="cuda") # TODO - upload checkpoint
+comparing_controlnet = SDXLControlNetPreEncodedControlnetCond.load("./controlnet_vae.safetensors", device="cuda") # TODO - upload checkpoint
 comparing_controlnet.to(torch.float16)
 
 def read_content(file_path: str) -> str:
@@ -40,13 +45,15 @@ def predict(dict, prompt="", negative_prompt="", guidance_scale=7.5, steps=20, s
     init_image = dict["image"].convert("RGB").resize((1024, 1024))
     mask = dict["mask"].convert("RGB").resize((1024, 1024))
     
-    output = pipe(prompt = prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength)
+    # output = pipe(prompt = prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength)
     output_controlnet_vae_encoding = gen_sdxl_simplified_interface(
-        prompt=prompt, negative_prompt=negative_prompt, images=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps),
+        prompts=prompt, negative_prompts=negative_prompt, images=init_image, masks=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps),
         text_encoder_one=pipe.text_encoder, text_encoder_two=pipe.text_encoder_2, unet=comparing_unet, vae=comparing_vae, controlnet=comparing_controlnet, device=device
     )
     
-    return output.images[0], output_controlnet_vae_encoding[0], gr.update(visible=True)
+    # return output.images[0], output_controlnet_vae_encoding[0], gr.update(visible=True)
+
+    return output_controlnet_vae_encoding[0], gr.update(visible=True)
 
 
 css = '''
@@ -108,16 +115,18 @@ with image_blocks as demo:
                             scheduler = gr.Dropdown(label="Schedulers", choices=schedulers, value="EulerDiscreteScheduler")
                         
                 with gr.Column():
-                    image_out = gr.Image(label="Output", elem_id="output-img", height=400)
-                    image_out_comparing = gr.Image(label="Output", elem_id="output-img-comparing", height=400)
+                    image_out = gr.Image(label="Output diffusers full finetune 0.1", elem_id="output-img", height=400)
+                    image_out_comparing = gr.Image(label="Output controlnet + vae", elem_id="output-img-comparing", height=400)
                     with gr.Group(elem_id="share-btn-container", visible=False) as share_btn_container:
                         community_icon = gr.HTML(community_icon_html)
                         loading_icon = gr.HTML(loading_icon_html)
                         share_button = gr.Button("Share to community", elem_id="share-btn",visible=True)
             
 
-    btn.click(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, image_out_comparing, share_btn_container], api_name='run')
-    prompt.submit(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, image_out_comparing, share_btn_container])
+    # btn.click(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, image_out_comparing, share_btn_container], api_name='run')
+    # prompt.submit(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, image_out_comparing, share_btn_container])
+    btn.click(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out_comparing, share_btn_container], api_name='run')
+    prompt.submit(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out_comparing, share_btn_container])
     share_button.click(None, [], [], _js=share_js)
 
     gr.Examples(

sdxl.py

@@ -388,9 +388,9 @@ def make_sample(d, proportion_empty_prompts, get_sdxl_conditioning_images=None):
 
     micro_conditioning = torch.tensor([original_width, original_height, c_top, c_left, 1024, 1024])
 
-    text_input_ids_one = sdxl_tokenize_one(text)
+    text_input_ids_one = sdxl_tokenize_one(text)[0]
 
-    text_input_ids_two = sdxl_tokenize_two(text)
+    text_input_ids_two = sdxl_tokenize_two(text)[0]
 
     image = image.convert("RGB")
 
@@ -517,7 +517,7 @@ def sdxl_tokenize_one(prompts):
         max_length=tokenizer_one.model_max_length,
         truncation=True,
         return_tensors="pt",
-    ).input_ids[0]
+    ).input_ids
 
 
 def sdxl_tokenize_two(prompts):
@@ -527,7 +527,7 @@ def sdxl_tokenize_two(prompts):
         max_length=tokenizer_one.model_max_length,
         truncation=True,
         return_tensors="pt",
-    ).input_ids[0]
+    ).input_ids
 
 
 def sdxl_text_conditioning(text_encoder_one, text_encoder_two, text_input_ids_one, text_input_ids_two):
@@ -667,7 +667,7 @@ def apply_padding(mask, coord):
 
 @torch.no_grad()
 def sdxl_diffusion_loop(
-    prompts,
+    prompts: List[str],
     unet,
     text_encoder_one,
     text_encoder_two,
@@ -683,8 +683,10 @@ def sdxl_diffusion_loop(
     negative_prompts=None,
     diffusion_loop=euler_ode_solver_diffusion_loop,
 ):
+    batch_size = len(prompts)
+
     if negative_prompts is None:
-        negative_prompts = [""] * len(prompts)
+        negative_prompts = [""] * batch_size
 
     prompts += negative_prompts
 
@@ -694,27 +696,30 @@ def sdxl_diffusion_loop(
         sdxl_tokenize_one(prompts).to(text_encoder_one.device),
         sdxl_tokenize_two(prompts).to(text_encoder_two.device),
     )
-
-    if x_T is None:
-        x_T = torch.randn((1, 4, 1024 // 8, 1024 // 8), dtype=torch.float32, device=unet.device, generator=generator)
-        x_T = x_T * ((sigmas.max() ** 2 + 1) ** 0.5)
+    encoder_hidden_states = encoder_hidden_states.to(unet.dtype)
+    pooled_encoder_hidden_states = pooled_encoder_hidden_states.to(unet.dtype)
 
     if sigmas is None:
         sigmas = make_sigmas(device=unet.device)
 
+    if x_T is None:
+        x_T = torch.randn((batch_size, 4, 1024 // 8, 1024 // 8), dtype=unet.dtype, device=unet.device, generator=generator)
+        x_T = x_T * ((sigmas.max() ** 2 + 1) ** 0.5)
+
     if timesteps is None:
         timesteps = torch.linspace(0, sigmas.numel(), 50, dtype=torch.long, device=unet.device)
 
     if micro_conditioning is None:
-        micro_conditioning = torch.tensor([1024, 1024, 0, 0, 1024, 1024], dtype=torch.long, device=unet.device)
+        micro_conditioning = torch.tensor([[1024, 1024, 0, 0, 1024, 1024]], dtype=torch.long, device=unet.device)
+        micro_conditioning = micro_conditioning.expand(batch_size, -1)
 
     if adapter is not None:
-        down_block_additional_residuals = adapter(images)
+        down_block_additional_residuals = adapter(images.to(dtype=adapter.dtype, device=adapter.device))
     else:
         down_block_additional_residuals = None
 
     if controlnet is not None:
-        controlnet_cond = images
+        controlnet_cond = images.to(dtype=controlnet.dtype, device=controlnet.device)
     else:
         controlnet_cond = None
 
@@ -756,21 +761,28 @@ def sdxl_eps_theta(
 
     if guidance_scale > 1.0:
         scaled_x_t = torch.concat([scaled_x_t, scaled_x_t])
+        micro_conditioning = torch.concat([micro_conditioning, micro_conditioning])
+        if controlnet_cond is not None:
+            controlnet_cond = torch.concat([controlnet_cond, controlnet_cond])
 
     if controlnet is not None:
         controlnet_out = controlnet(
             x_t=scaled_x_t,
             t=t,
-            encoder_hidden_states=encoder_hidden_states,
-            micro_conditioning=micro_conditioning,
-            pooled_encoder_hidden_states=pooled_encoder_hidden_states,
+            encoder_hidden_states=encoder_hidden_states.to(controlnet.dtype),
+            micro_conditioning=micro_conditioning.to(controlnet.dtype),
+            pooled_encoder_hidden_states=pooled_encoder_hidden_states.to(controlnet.dtype),
             controlnet_cond=controlnet_cond,
         )
 
-        down_block_additional_residuals = controlnet_out["down_block_res_samples"]
-        mid_block_additional_residual = controlnet_out["mid_block_res_sample"]
+        down_block_additional_residuals = [x.to(unet.dtype) for x in controlnet_out["down_block_res_samples"]]
+        mid_block_additional_residual = controlnet_out["mid_block_res_sample"].to(unet.dtype)
         add_to_down_block_inputs = controlnet_out.get("add_to_down_block_inputs", None)
+        if add_to_down_block_inputs is not None:
+            add_to_down_block_inputs = [x.to(unet.dtype) for x in add_to_down_block_inputs]
         add_to_output = controlnet_out.get("add_to_output", None)
+        if add_to_output is not None:
+            add_to_output = add_to_output.to(unet.dtype)
     else:
         mid_block_additional_residual = None
         add_to_down_block_inputs = None
@@ -795,20 +807,24 @@ def sdxl_eps_theta(
 
     return eps_hat
 
+
 known_negative_prompt = "text, watermark, low-quality, signature, moiré pattern, downsampling, aliasing, distorted, blurry, glossy, blur, jpeg artifacts, compression artifacts, poorly drawn, low-resolution, bad, distortion, twisted, excessive, exaggerated pose, exaggerated limbs, grainy, symmetrical, duplicate, error, pattern, beginner, pixelated, fake, hyper, glitch, overexposed, high-contrast, bad-contrast"
 
+
+# TODO probably just combine with sdxl_diffusion_loop
 def gen_sdxl_simplified_interface(
-    prompt:str,
-    negative_prompt: Optional[str] = None,
-    controlnet_checkpoint: Optional[str]=None,
-    controlnet: Optional[Literal["SDXLControlNet", "SDXLContolNetFull", "SDXLControlNetPreEncodedControlnetCond"]]=None,
-    adapter_checkpoint: Optional[str]=None,
+    prompts: Union[str, List[str]],
+    negative_prompts: Optional[Union[str, List[str]]] = None,
+    controlnet_checkpoint: Optional[str] = None,
+    controlnet: Optional[Literal["SDXLControlNet", "SDXLContolNetFull", "SDXLControlNetPreEncodedControlnetCond"]] = None,
+    adapter_checkpoint: Optional[str] = None,
     num_inference_steps=50,
     images=None,
     masks=None,
-    apply_conditioning: Optional[Literal["canny"]]=None,
-    num_images: int=1,
-    device: Optional[str]=None,
+    apply_conditioning: Optional[Literal["canny"]] = None,
+    num_images: int = 1,
+    guidance_scale=5.0,
+    device: Optional[str] = None,
     text_encoder_one=None,
     text_encoder_two=None,
     unet=None,
@@ -886,22 +902,23 @@ def gen_sdxl_simplified_interface(
                 mask = masks[image_idx]
                 if isinstance(mask, str):
                     mask = Image.open(mask)
-                    mask = mask.convert("L")
-                    mask = mask.resize((1024, 1024))
                 elif isinstance(mask, Image.Image):
                     ...
                 else:
                     assert False
+                mask = mask.convert("L")
+                mask = mask.resize((1024, 1024))
                 mask = TF.to_tensor(mask)
 
-                if controlnet == "SDXLControlNetPreEncodedControlnetCond":
+                if isinstance(controlnet, SDXLControlNetPreEncodedControlnetCond):
                     image = image * (mask < 0.5)
-                    image = TF.normalized(image, [0.5], [0.5])
-                    image = vae.encode(image)
-                    mask = TF.resize(mask, (1024 // 8, 1024 // 8))
-                    image = torch.concat((image, mask))
+                    image = TF.normalize(image, [0.5], [0.5])
+                    image = vae.encode(image[None, :, :, :].to(dtype=vae.dtype, device=vae.device)).to(dtype=unet.dtype, device=unet.device)
+                    mask = TF.resize(mask, (1024 // 8, 1024 // 8))[None, :, :, :].to(dtype=image.dtype, device=image.device)
+                    image = torch.concat((image, mask), dim=1)
                 else:
-                    image = image * (mask < 0.5) + -1.0 * (mask >= 0.5)
+                    image = (image * (mask < 0.5) + -1.0 * (mask >= 0.5)).to(dtype=unet.dtype, device=unet.device)
+                    image = image[None, :, :, :]
 
             images_.append(image)
 
@@ -909,9 +926,24 @@ def gen_sdxl_simplified_interface(
     else:
         images_ = None
 
+    if isinstance(prompts, str):
+        prompts = [prompts]
+    prompts_ = []
+    for prompt in prompts:
+        prompts_ += [prompt] * num_images
+
+    if negative_prompts is not None:
+        if isinstance(negative_prompts, str):
+            negative_prompts = [negative_prompts]
+        negative_prompts_ = []
+        for negative_prompt in negative_prompts:
+            negative_prompts_ += [negative_prompt] * num_images
+    else:
+        negative_prompts_ = None
+
     x_0 = sdxl_diffusion_loop(
-        prompts=[prompt] * num_images,
-        negative_prompts=[negative_prompt] * num_images,
+        prompts=prompts_,
+        negative_prompts=negative_prompts_,
         unet=unet,
         text_encoder_one=text_encoder_one,
         text_encoder_two=text_encoder_two,
@@ -920,9 +952,10 @@ def gen_sdxl_simplified_interface(
         controlnet=controlnet,
         adapter=adapter,
         images=images_,
+        guidance_scale=guidance_scale,
     )
 
-    x_0 = vae.decode(x_0)
+    x_0 = vae.decode(x_0.to(vae.dtype))
     x_0 = vae.output_tensor_to_pil(x_0)
 
     return x_0

sdxl_models.py

@@ -26,7 +26,8 @@ class ModelUtils:
 
         load_from = [load_from]
 
-        load_from += overrides
+        if overrides is not None:
+            load_from += overrides
 
         state_dict = {}
 
@@ -79,7 +80,7 @@ class SDXLVae(nn.Module, ModelUtils):
 
             # 512 -> 512
             mid_block=nn.ModuleDict(dict(
-                attentions=nn.ModuleList([Attention(512, 512, qkv_bias=True)]),
+                attentions=nn.ModuleList([VaeMidBlockAttention(512)]),
                 resnets=nn.ModuleList([ResnetBlock2D(512, 512, eps=1e-6), ResnetBlock2D(512, 512, eps=1e-6)]),
             )),
 
@@ -95,7 +96,7 @@ class SDXLVae(nn.Module, ModelUtils):
         # 8 -> 4 from sampling mean and std
 
         # 4 -> 4
-        self.post_quant_conv = nn.Conv2d(4, 4, 1)
+        self.post_quant_conv = nn.Conv2d(4, 4, kernel_size=1)
 
         self.decoder = nn.ModuleDict(dict(
             # 4 -> 512
@@ -103,7 +104,7 @@ class SDXLVae(nn.Module, ModelUtils):
 
             # 512 -> 512
             mid_block=nn.ModuleDict(dict(
-                attentions=nn.ModuleList([Attention(512, 512, qkv_bias=True)]),
+                attentions=nn.ModuleList([VaeMidBlockAttention(512)]),
                 resnets=nn.ModuleList([ResnetBlock2D(512, 512, eps=1e-6), ResnetBlock2D(512, 512, eps=1e-6)]),
             )),
 
@@ -179,15 +180,18 @@ class SDXLVae(nn.Module, ModelUtils):
 
         h = self.post_quant_conv(h)
 
+        h = self.decoder["conv_in"](h)
+
         h = self.decoder["mid_block"]["resnets"][0](h)
         h = self.decoder["mid_block"]["attentions"][0](h)
         h = self.decoder["mid_block"]["resnets"][1](h)
 
-        for up_block in self.encoder["up_blocks"]:
+        for up_block in self.decoder["up_blocks"]:
             for resnet in up_block["resnets"]:
                 h = resnet(h)
 
             if "upsamplers" in up_block:
+                h = F.interpolate(h, scale_factor=2.0, mode="nearest")
                 h = up_block["upsamplers"][0]["conv"](h)
 
         h = self.decoder["conv_norm_out"](h)
@@ -208,9 +212,7 @@ class SDXLVae(nn.Module, ModelUtils):
 
     @classmethod
     def output_tensor_to_pil(self, x_pred):
-        x_pred = ((x_pred * 0.5 + 0.5).clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 1)
-
-        x_pred = x_pred.permute(0, 2, 3, 1).cpu().numpy()
+        x_pred = ((x_pred * 0.5 + 0.5).clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy()
 
         x_pred = [Image.fromarray(x) for x in x_pred]
 
@@ -1323,42 +1325,83 @@ class TransformerDecoderBlock(nn.Module):
 
 
 class Attention(nn.Module):
-    def __init__(self, channels, encoder_hidden_states_dim, qkv_bias=False):
+    def __init__(self, channels, encoder_hidden_states_dim):
         super().__init__()
-        self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
-        self.to_k = nn.Linear(encoder_hidden_states_dim, channels, bias=qkv_bias)
-        self.to_v = nn.Linear(encoder_hidden_states_dim, channels, bias=qkv_bias)
+        self.to_q = nn.Linear(channels, channels, bias=False)
+        self.to_k = nn.Linear(encoder_hidden_states_dim, channels, bias=False)
+        self.to_v = nn.Linear(encoder_hidden_states_dim, channels, bias=False)
         self.to_out = nn.Sequential(nn.Linear(channels, channels), nn.Dropout(0.0))
 
     def forward(self, hidden_states, encoder_hidden_states=None):
-        batch_size, q_seq_len, channels = hidden_states.shape
-        head_dim = 64
+        input_ndim = hidden_states.ndim
 
-        if encoder_hidden_states is not None:
-            kv = encoder_hidden_states
-        else:
-            kv = hidden_states
+        if input_ndim == 4:
+            batch_size, channels, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)
+
+        hidden_states = attention(self.to_q, self.to_k, self.to_v, self.to_out, hidden_states, encoder_hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(1, 2).view(batch_size, channels, height, width)
 
-        kv_seq_len = kv.shape[1]
+        return hidden_states
+
+
+class VaeMidBlockAttention(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.group_norm = nn.GroupNorm(32, channels, eps=1e-06)
+        self.to_q = nn.Linear(channels, channels, bias=True)
+        self.to_k = nn.Linear(channels, channels, bias=True)
+        self.to_v = nn.Linear(channels, channels, bias=True)
+        self.to_out = nn.Sequential(nn.Linear(channels, channels), nn.Dropout(0.0))
 
-        query = self.to_q(hidden_states)
-        key = self.to_k(kv)
-        value = self.to_v(kv)
+    def forward(self, hidden_states):
+        input_ndim = hidden_states.ndim
 
-        query = query.reshape(batch_size, q_seq_len, channels // head_dim, head_dim).contiguous()
-        key = key.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).contiguous()
-        value = value.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).contiguous()
+        if input_ndim == 4:
+            batch_size, channels, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)
 
-        hidden_states = xformers.ops.memory_efficient_attention(query, key, value)
+        hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
 
-        hidden_states = hidden_states.to(query.dtype)
-        hidden_states = hidden_states.reshape(batch_size, q_seq_len, channels).contiguous()
+        hidden_states = attention(self.to_q, self.to_k, self.to_v, self.to_out, hidden_states)
 
-        hidden_states = self.to_out(hidden_states)
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(1, 2).view(batch_size, channels, height, width)
 
         return hidden_states
 
 
+def attention(to_q, to_k, to_v, to_out, hidden_states, encoder_hidden_states=None):
+    batch_size, q_seq_len, channels = hidden_states.shape
+    head_dim = 64
+
+    if encoder_hidden_states is not None:
+        kv = encoder_hidden_states
+    else:
+        kv = hidden_states
+
+    kv_seq_len = kv.shape[1]
+
+    query = to_q(hidden_states)
+    key = to_k(kv)
+    value = to_v(kv)
+
+    query = query.reshape(batch_size, q_seq_len, channels // head_dim, head_dim).contiguous()
+    key = key.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).contiguous()
+    value = value.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).contiguous()
+
+    hidden_states = xformers.ops.memory_efficient_attention(query, key, value)
+
+    hidden_states = hidden_states.to(query.dtype)
+    hidden_states = hidden_states.reshape(batch_size, q_seq_len, channels).contiguous()
+
+    hidden_states = to_out(hidden_states)
+
+    return hidden_states
+
+
 class GEGLU(nn.Module):
     def __init__(self, dim_in: int, dim_out: int):
         super().__init__()