make karras sigmas and sde_variance work together

Author: limiteinductive

diffusers_ddim.py

@@ -0,0 +1,37 @@
+import torch
+from diffusers import StableDiffusionPipeline
+
+# from diffusers.schedulers.scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
+from diffusers.schedulers.scheduling_ddim import DDIMScheduler
+
+from refiners.fluxion.utils import manual_seed
+
+# diffusers_solver = DDIMScheduler.from_config( # type: ignore
+#     {
+#         "_class_name": "PNDMScheduler",
+#         "_diffusers_version": "0.6.0",
+#         "beta_end": 0.012,
+#         "beta_schedule": "scaled_linear",
+#         "beta_start": 0.00085,
+#         "num_train_timesteps": 1000,
+#         "set_alpha_to_one": False,
+#         "skip_prk_steps": True,
+#         "steps_offset": 1,
+#         "trained_betas": None,
+#         "clip_sample": False,
+#         # "use_karras_sigmas": True,
+#         # "algorithm": "dpmsolver++",
+#     }
+# )
+
+
+model_id = "botp/stable-diffusion-v1-5"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, scheduler=diffusers_solver)
+pipe = pipe.to("cuda:1")
+
+prompt = "a cute cat, detailed high-quality professional image"
+negative_prompt = "lowres, bad anatomy, bad hands, cropped, worst quality"
+manual_seed(2)
+image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
+
+image.save("diffusers-ddim.png")

diffusers_karrs.py

@@ -0,0 +1,22 @@
+import torch
+from diffusers import StableDiffusionPipeline
+from diffusers.schedulers.scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
+
+from refiners.fluxion.utils import manual_seed
+
+model_id = "botp/stable-diffusion-v1-5"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+pipe = pipe.to("cuda:1")
+
+
+config = {**pipe.scheduler.config}
+config["use_karras_sigmas"] = True
+config["algorithm_type"] = "sde-dpmsolver++"
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(config)
+
+prompt = "a cute cat, detailed high-quality professional image"
+negative_prompt = "lowres, bad anatomy, bad hands, cropped, worst quality"
+manual_seed(2)
+image = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=18, guidance_scale=7.5).images[0]
+
+image.save("diffusers-sde-dpm-karras.png")

old-script.py

@@ -0,0 +1,30 @@
+import torch
+from PIL import Image
+
+from diffusers import DPMSolverMultistepScheduler
+from diffusers import StableDiffusionPipeline
+
+weights_path = "weights"
+device = torch.device("cuda:1")
+n_steps = 30
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    "botp/stable-diffusion-v1-5",
+    torch_dtype=torch.float32,
+    safety_checker=None,
+).to(device)
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+
+prompt = "a cute cat, detailed high-quality professional image"
+negative_prompt = "lowres, bad anatomy, bad hands, cropped, worst quality"
+
+
+torch.manual_seed(2)
+output = pipe(
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    num_inference_steps=n_steps,
+    guidance_scale=7.5,
+)
+
+output.images[0].save("output-diffusers-diffusion.png")

run_ic_light.py

@@ -0,0 +1,49 @@
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+from refiners.fluxion.utils import load_from_safetensors, no_grad
+from refiners.foundationals.clip.text_encoder import CLIPTextEncoderL
+from refiners.foundationals.latent_diffusion.stable_diffusion_1.ic_light import ICLightBackground
+from refiners.foundationals.latent_diffusion.stable_diffusion_1.model import SD1Autoencoder
+from refiners.foundationals.latent_diffusion.stable_diffusion_1.unet import SD1UNet
+
+no_grad().__enter__()
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+sd = ICLightBackground(
+    patch_weights=load_from_safetensors("iclight_sd15_fbc-refiners.safetensors", device=device),
+    unet=SD1UNet(in_channels=4, device=device).load_from_safetensors("realistic-vision-v51-unet.safetensors"),
+    clip_text_encoder=CLIPTextEncoderL(device=device).load_from_safetensors(
+        "realistic-vision-v51-text_encoder.safetensors"
+    ),
+    lda=SD1Autoencoder(device=device).load_from_safetensors("realistic-vision-v51-autoencoder.safetensors"),
+    device=device,
+)
+
+prompt = "porcelaine mug, 4k high quality, soft lighting, high-quality professional image"
+negative_prompt = "lowres, bad anatomy, bad hands, cropped, worst quality"
+clip_text_embedding = sd.compute_clip_text_embedding(text=prompt, negative_text=negative_prompt)
+
+data_test = Path.home() / "data_test"
+uid = "ic_out"
+foreground_image = Image.open(f"{uid}.png").convert("RGB")
+mask = Image.open(f"{uid}.mask.png").resize(foreground_image.size).convert("L").point(lambda x: 255 if x > 0 else 0)
+
+
+sd.set_foreground_condition(foreground_image, mask=mask, use_rescaled_image=True)
+sd.set_background_condition(foreground_image)
+
+x = torch.randn(1, 4, foreground_image.height // 8, foreground_image.width // 8, device=device)
+
+for step in sd.steps:
+    x = sd(
+        x,
+        step=step,
+        clip_text_embedding=clip_text_embedding,
+        condition_scale=3,
+    )
+predicted_image = sd.lda.latents_to_image(x)
+
+predicted_image.save("ic-light-output.png")

src/refiners/foundationals/latent_diffusion/flux/__init__.py

@@ -0,0 +1,92 @@
+from typing import Annotated, NamedTuple, cast
+
+import torch
+
+from refiners.fluxion import layers as fl
+
+
+class FluxParams(NamedTuple):
+    in_channels: int
+    vec_in_dim: int
+    context_in_dim: int
+    hidden_size: int
+    mlp_ratio: float
+    num_heads: int
+    depth: int
+    depth_single_blocks: int
+    axes_dim: list[int]
+    theta: int
+    qkv_bias: bool
+    guidance_embed: bool
+
+
+class DoubleStreamBlocks(fl.Chain):
+    pass
+
+
+class SingleStreamBlocks(fl.Chain):
+    pass
+
+
+class PositionalEncoding(fl.Passthrough):
+    pass
+
+
+class TimestepEmbedding(fl.Passthrough):
+    pass
+
+
+class TextEmbedding(fl.Passthrough):
+    pass
+
+
+class DoubleStreamAttention(fl.Chain):
+    pass
+
+
+class SingleStreamAttention(fl.Chain):
+    pass
+
+
+class LastLayer(fl.Chain):
+    def __init__(self, params: FluxParams) -> None:
+        super().__init__(fl.LayerNorm(params.hidden_size), fl.Linear(params.hidden_size, params.in_channels))
+
+
+class Flux(fl.Sum):
+    def __init__(
+        self,
+        params: FluxParams,
+    ) -> None:
+        super().__init__(
+            TimestepEmbedding(),
+            PositionalEncoding(),
+            TextEmbedding(),
+            DoubleStreamBlocks(DoubleStreamAttention() for _ in range(params.depth)),
+            SingleStreamBlocks(SingleStreamAttention() for _ in range(params.depth_single_blocks)),
+        )
+
+
+if __name__ == "__main__":
+    flux = Flux(
+        params=FluxParams(
+            in_channels=3,
+            vec_in_dim=32,
+            context_in_dim=64,
+            hidden_size=256,
+            mlp_ratio=4.0,
+            num_heads=8,
+            depth=2,
+            depth_single_blocks=1,
+            axes_dim=[32, 32],
+            theta=1024,
+            qkv_bias=True,
+            guidance_embed=True,
+        ),
+    )
+    from typing import get_type_hints
+
+    print(repr(flux))
+    image_in = flux.layer(0)
+    print(image_in)
+    print(get_type_hints(image_in))