import os, math, numpy as np from PIL import Image import torch from diffusers import AutoencoderKL # ---- TensorFlow / Keras (kept, as requested) ---- import tensorflow as tf from tensorflow import keras from keras import layers # ---- Gradio UI ---- import gradio as gr # ---- CLIP (OpenAI repo) ---- import clip # ---- Hugging Face Hub (to fetch your .h5) ---- from huggingface_hub import hf_hub_download # ---------------------------- # Hub model location (YOUR PUBLIC REPO) # ---------------------------- MODEL_REPO_ID = "Manjeet9812/LDM_trained_model" MODEL_FILENAME = "text_to_image.h5" # latent -> pixel scale for SD VAE STD_LATENT_DEFAULT = 10.5 LATENT_H = 16 # your code uses 16x16x4 latents LATENT_W = 16 LATENT_C = 4 # ---------------------------- # Utils (kept from your code) # ---------------------------- def attention(qkv): q, k, v = qkv vector = tf.matmul(k, q, transpose_b=True) score = tf.nn.softmax(vector) o = tf.matmul(score, v) return o def spatial_attention(img): filters = img.shape[3] orig_shape = (img.shape[1], img.shape[2], img.shape[3]) img = layers.BatchNormalization()(img) q = layers.Conv2D(filters // 8, 1, padding="same")(img) k = layers.Conv2D(filters // 8, 1, padding="same")(img) v = layers.Conv2D(filters, 1, padding="same")(img) k = layers.Reshape((k.shape[1] * k.shape[2], k.shape[3]))(k) q = layers.Reshape((q.shape[1] * q.shape[2], q.shape[3]))(q) v = layers.Reshape((v.shape[1] * v.shape[2], v.shape[3]))(v) img = layers.Lambda(attention)([q, k, v]) img = layers.Reshape(orig_shape)(img) img = layers.Conv2D(filters, 1, padding="same")(img) img = layers.BatchNormalization()(img) return img def cross_attention(img, text): filters = img.shape[3] orig_shape = (img.shape[1], img.shape[2], img.shape[3]) img = layers.BatchNormalization()(img) text = layers.BatchNormalization()(text) q = layers.Conv2D(filters // 8, 1, padding="same")(text) k = layers.Conv2D(filters // 8, 1, padding="same")(img) v = layers.Conv2D(filters, 1, padding="same")(text) q = layers.Reshape((q.shape[1] * q.shape[2], q.shape[3]))(q) k = layers.Reshape((k.shape[1] * k.shape[2], k.shape[3]))(k) v = layers.Reshape((v.shape[1] * v.shape[2], v.shape[3]))(v) img = layers.Lambda(attention)([q, k, v]) img = layers.Reshape(orig_shape)(img) img = layers.Conv2D(filters, 1, padding="same")(img) img = layers.BatchNormalization()(img) return img def sinusoidal_embedding(x): embedding_min_frequency = 1.0 embedding_max_frequency = 1000.0 embedding_dims = 32 frequencies = tf.exp( tf.linspace( tf.math.log(embedding_min_frequency), tf.math.log(embedding_max_frequency), embedding_dims // 2, ) ) angular_speeds = 2.0 * math.pi * frequencies embeddings = tf.concat( [tf.sin(angular_speeds * x), tf.cos(angular_speeds * x)], axis=3 ) return embeddings def dynamic_thresholding(img, perc=99.5): s = np.percentile(np.abs(img.ravel()), perc) s = np.max([s, 1]) img = img.clip(-s, s) / s return img class Diffuser: def __init__(self, denoiser, class_guidance, diffusion_steps, perc_thresholding=99.5, batch_size=64): self.denoiser = denoiser self.class_guidance = class_guidance self.diffusion_steps = diffusion_steps self.noise_levels = 1 - np.power(np.arange(0.0001, 0.99, 1 / self.diffusion_steps), 1 / 3) self.noise_levels[-1] = 0.01 self.perc_thresholding = perc_thresholding self.batch_size = batch_size def predict_x_zero(self, x_t, label, noise_level): num_imgs = len(x_t) label_empty_ohe = np.zeros(shape=label.shape) noise_in = np.array([noise_level] * num_imgs)[:, None, None, None] nn_inputs = [np.vstack([x_t, x_t]), np.vstack([noise_in, noise_in]), np.vstack([label, label_empty_ohe])] x0_pred = self.denoiser.predict(nn_inputs, batch_size=self.batch_size, verbose=0) x0_pred_label = x0_pred[:num_imgs] x0_pred_no_label = x0_pred[num_imgs:] x0_pred = self.class_guidance * x0_pred_label + (1 - self.class_guidance) * x0_pred_no_label x0_pred = dynamic_thresholding(x0_pred, perc=self.perc_thresholding) return x0_pred def reverse_diffusion(self, seeds, label, show_img=False): new_img = seeds for _ in range(len(self.noise_levels) - 1): curr_noise, next_noise = self.noise_levels[_], self.noise_levels[_ + 1] x0_pred = self.predict_x_zero(new_img, label, curr_noise) new_img = ((curr_noise - next_noise) * x0_pred + next_noise * new_img) / curr_noise return x0_pred def decode_latents(latent_arr, std_latent=STD_LATENT_DEFAULT, batch_size=32, vae=None, torch_device="cpu"): assert vae is not None decoded_imgs = [] N = latent_arr.shape[0] with torch.no_grad(): for start in range(0, N, batch_size): end = min(start + batch_size, N) encoded = torch.from_numpy(latent_arr[start:end] * float(std_latent)).permute(0, 3, 1, 2).to(torch_device, dtype=torch.float32) decoded = vae.decode(encoded).sample decoded_imgs.append(decoded.permute(0, 2, 3, 1).cpu().numpy()) return np.concatenate(decoded_imgs, axis=0) def to_pil_list(imgs_np): # imgs in [-1, 1] -> [0,255] uint8 imgs = np.clip((imgs_np + 1.0) * 127.5, 0, 255).astype("uint8") out = [] for i in range(imgs.shape[0]): out.append(Image.fromarray(imgs[i])) return out # ---------------------------- # Load models (on import) # ---------------------------- torch_device = "cuda" if torch.cuda.is_available() else "cpu" # SD VAE vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema").to(torch_device) vae.eval() # Keras custom objects custom_objects = {"sinusoidal_embedding": sinusoidal_embedding, "attention": attention} # Download your UNet .h5 from the Hub and load it try: LOCAL_MODEL_PATH = hf_hub_download(repo_id=MODEL_REPO_ID, filename=MODEL_FILENAME) except Exception as e: raise RuntimeError( f"Could not download {MODEL_FILENAME} from {MODEL_REPO_ID}. " f"Error: {e}" ) unet = keras.models.load_model(LOCAL_MODEL_PATH, custom_objects=custom_objects, compile=False) # CLIP (kept as in your test code) clip_model, _ = clip.load("ViT-B/32", device=torch_device) # ---------------------------- # Gradio inference wrapper # ---------------------------- def generate( prompt: str, num_images: int = 4, diffusion_steps: int = 100, class_guidance: float = 6.0, perc_thresholding: float = 99.75, std_latent: float = STD_LATENT_DEFAULT, seed: int | None = None, batch_size: int = 32, ): if not prompt or prompt.strip() == "": return [] # Seed if seed is not None and int(seed) != 0: np.random.seed(int(seed)) # Text -> CLIP embedding (repeat for batch) with torch.no_grad(): text_tokens = clip.tokenize(prompt, truncate=False).to(torch_device) text_encoding = clip_model.encode_text(text_tokens) # (1, D) text_encoding = text_encoding.detach().float() text_np = np.vstack([text_encoding.cpu().numpy()] * int(num_images)) # (N, D) # Random latent seeds (NHWC) rand_latents = np.random.normal(0, 1, (int(num_images), LATENT_H, LATENT_W, LATENT_C)).astype("float32") # Run your diffuser (with your UNet) diffuser = Diffuser( denoiser=unet, class_guidance=float(class_guidance), diffusion_steps=int(diffusion_steps), perc_thresholding=float(perc_thresholding), batch_size=int(batch_size), ) # TF GPU is usually not present on Spaces → keep UNet predict on CPU with tf.device("/CPU:0"): imgs_latent = diffuser.reverse_diffusion(rand_latents, text_np) # Decode with SD VAE with torch.no_grad(): imgs = decode_latents( imgs_latent, std_latent=float(std_latent), batch_size=int(batch_size), vae=vae, torch_device=torch_device ) # Convert to PIL return to_pil_list(imgs) # ---------------------------- # Gradio UI # ---------------------------- with gr.Blocks(title="LDM Host (Keras UNet + SD VAE + CLIP)") as demo: gr.Markdown("## LDM Host — Text → Image (kept same logic)") with gr.Row(): with gr.Column(scale=1): prompt = gr.Textbox(label="Text prompt", value="beautiful landscape with river and flowers") num_images = gr.Slider(1, 8, value=4, step=1, label="Number of images") diffusion_steps = gr.Slider(10, 200, value=100, step=1, label="Diffusion steps") class_guidance = gr.Slider(0.0, 10.0, value=6.0, step=0.25, label="Classifier-free guidance (blend)") perc_thr = gr.Slider(90.0, 99.95, value=99.75, step=0.05, label="Dynamic thresholding percentile") std_latent = gr.Slider(1.0, 20.0, value=STD_LATENT_DEFAULT, step=0.1, label="Latent std scale (SD VAE)") batch_size = gr.Slider(1, 64, value=32, step=1, label="Batch size (UNet predict)") seed = gr.Number(label="Seed (0 = random)", value=0, precision=0) run = gr.Button("Generate") with gr.Column(scale=1): gallery = gr.Gallery(label="Outputs", columns=2, rows=2, height=560) run.click( fn=generate, inputs=[prompt, num_images, diffusion_steps, class_guidance, perc_thr, std_latent, seed, batch_size], outputs=[gallery], api_name="generate" ) if __name__ == "__main__": demo.launch()