Add SD3LatentFormat and CFGDenoiser

keras_cv/src/models/stable_diffusion_v3/cfg_denoiser.py

@@ -0,0 +1,9 @@
+from keras import ops
+
+
+class CFGDenoiser:
+    def __call__(self, batched, cond_scale):
+        # `batched` is the outputs from `BaseModel.apply_model`
+        pos_out, neg_out = ops.split(batched, 2, axis=0)
+        scaled = neg_out + (pos_out - neg_out) * cond_scale
+        return scaled

keras_cv/src/models/stable_diffusion_v3/sd3_latent_format.py

@@ -0,0 +1,15 @@
+class SD3LatentFormat:
+    """Latents are slightly shifted from center.
+
+    This class must be called after VAE Decode to correct for the shift.
+    """
+
+    def __init__(self):
+        self.scale_factor = 1.5305
+        self.shift_factor = 0.0609
+
+    def process_in(self, latent):
+        return (latent - self.shift_factor) * self.scale_factor
+
+    def process_out(self, latent):
+        return (latent / self.scale_factor) + self.shift_factor