import gradio as gr import numpy as np import torch import re import os from diffusers import UNet2DConditionModel, DDPMScheduler, AutoencoderKL from transformers import CLIPTokenizer, CLIPTextModel, CLIPModel from huggingface_hub import snapshot_download import torch.nn.functional as F from torchvision.transforms import transforms from PIL import Image # ── Load models ONCE at startup ────────────────────────────────────────────── device = "cuda" if torch.cuda.is_available() else "cpu" torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 LATENT_SIZE = 16 print("Downloading model weights...") model_path = snapshot_download( repo_id="vish26/latent-diffusion-model-128x128-batch8-lr-1e-5", repo_type="model" ) model_path = os.path.join(model_path, "epoch_7") print("Loading models...") unet = UNet2DConditionModel.from_pretrained( os.path.join(model_path, "unet"), use_safetensors=True, torch_dtype=torch_dtype, low_cpu_mem_usage=True ).to(device) text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32", torch_dtype=torch_dtype).to(device) tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32") vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch_dtype).to(device) clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32", torch_dtype=torch_dtype).to(device) scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="linear") # Load training state checkpoint = torch.load(os.path.join(model_path, "training_state.pth"), map_location=device) unet.load_state_dict(checkpoint['model_state_dict']) print(f"Model loaded! Last loss: {checkpoint['loss']}") unet.eval() text_encoder.eval() vae.eval() clip.eval() print("All models ready!") def is_valid_prompt(prompt): # Empty or spaces if not prompt or prompt.strip() == "": return False # Must contain at least one letter or number if not re.search(r"[a-zA-Z0-9]", prompt): return False return True # ── Generation function ─────────────────────────────────────────────────────── def generate_image(prompt, num_inference_steps=500, guidance_scale=7.5): if not is_valid_prompt(prompt): raise gr.Error("Please enter a valid prompt (not empty or special characters only).") # Tokenize text_input = tokenizer( prompt, padding="max_length", max_length=77, truncation=True, return_tensors="pt" ).to(device) uncond_input = tokenizer( [""], padding="max_length", max_length=77, truncation=True, return_tensors="pt" ).to(device) with torch.no_grad(): text_embeddings = text_encoder(text_input.input_ids)[0] uncond_embeddings = text_encoder(uncond_input.input_ids)[0] text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) # Init latents latents = torch.randn((1, 4, LATENT_SIZE, LATENT_SIZE)).to(device) if torch_dtype == torch.float16: latents = latents.half() latents = latents * scheduler.init_noise_sigma # Denoising loop scheduler.set_timesteps(num_inference_steps) for t in scheduler.timesteps: latent_model_input = torch.cat([latents] * 2) latent_model_input = scheduler.scale_model_input(latent_model_input, t) with torch.no_grad(): noise_pred = unet( latent_model_input, t, encoder_hidden_states=text_embeddings ).sample noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) latents = scheduler.step(noise_pred, t, latents).prev_sample # Decode latents = 1 / 0.18215 * latents with torch.no_grad(): image = vae.decode(latents).sample image = (image / 2 + 0.5).clamp(0, 1) image = image.detach().cpu().float() # back to float32 for PIL image = image.permute(0, 2, 3, 1).squeeze(0).numpy() image = (image * 255).round().astype("uint8") image = Image.fromarray(image) # CLIP Score # transform = transforms.Compose([ # transforms.Resize((224, 224)), # transforms.ToTensor(), # transforms.Normalize([0.5], [0.5]) # ]) # processed_image = transform(image).unsqueeze(0).to(device) # if torch_dtype == torch.float16: # processed_image = processed_image.half() # with torch.no_grad(): # # ✅ Extract .image_embeds and .text_embeds from the output object # # image_features = clip.get_image_features(pixel_values=processed_image) # # image_features = F.normalize(image_features, dim=-1) # # text_features = clip.get_text_features(input_ids=text_input.input_ids) # # text_features = F.normalize(text_features, dim=-1) # image_features = clip.get_image_features(pixel_values=processed_image) # if hasattr(image_features, 'image_embeds'): # image_features = image_features.image_embeds # image_features = F.normalize(image_features, dim=-1) # text_features = clip.get_text_features(input_ids=text_input.input_ids) # if hasattr(text_features, 'text_embeds'): # text_features = text_features.text_embeds # text_features = F.normalize(text_features, dim=-1) # clip_score = (image_features * text_features).sum(dim=-1).item() return image # ── Gradio UI ───────────────────────────────────────────────────────────────── with gr.Blocks() as demo: gr.Markdown("# 🖼️ Latent Diffusion Model — Text to Image") with gr.Row(): with gr.Column(): prompt = gr.Text(label="Prompt", placeholder="e.g. people walking on street") steps = gr.Slider(label="Inference Steps", minimum=100, maximum=1000, step=50, value=500) guidance = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=15.0, step=0.5, value=7.5) generate_button = gr.Button("Generate", variant="primary") with gr.Column(): output_image = gr.Image(label="Generated Image") # clip_score = gr.Text(label="CLIP Score") generate_button.click( fn=generate_image, inputs=[prompt, steps, guidance], # outputs=[output_image, clip_score] outputs=[output_image] ) gr.Examples( examples=[ ["people walking on street", 500, 7.5], ["a dog playing in the park", 500, 7.5], ["sunset over mountains", 500, 7.5], ], inputs=[prompt, steps, guidance] ) if __name__ == "__main__": demo.launch()