""" DiffusionPen: Hindi Handwriting Generation Demo Inference-focused Gradio application with CANINE text encoding """ import gradio as gr import torch import numpy as np from PIL import Image from unet import UNetModel from transformers import CanineTokenizer, CanineModel from pathlib import Path class DiffusionPenDemo: """ Hindi Handwriting Generation Demo using DiffusionPen UNet Features: - CANINE text encoder for character-level Hindi encoding - 339 different writer styles - Configurable diffusion steps and guidance - GPU/CPU automatic detection - Checkpoint loading support """ def __init__(self, checkpoint_path=None, device=None): self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu') self.checkpoint_path = checkpoint_path self.model = None self.text_encoder = None self.tokenizer = None self.checkpoint_loaded = False self.load_models() def load_models(self): """Load UNet model and CANINE text encoder""" try: print(f"\n{'='*60}") print(f"🔧 DiffusionPen Initialization") print(f"{'='*60}") print(f"📱 Device: {self.device.upper()}") # Load CANINE text encoder print("\n📝 Loading CANINE text encoder...") self.tokenizer = CanineTokenizer.from_pretrained('google/canine-s') self.text_encoder = CanineModel.from_pretrained('google/canine-s').to(self.device) self.text_encoder.eval() print(" ✓ CANINE loaded (768-dim embeddings)") # Initialize UNet model print("\n🧠 Initializing UNet model...") class Args: interpolation = False mix_rate = 0.5 self.model = UNetModel( image_size=64, in_channels=1, model_channels=128, out_channels=1, num_res_blocks=2, attention_resolutions=[16, 8], dropout=0.1, channel_mult=(1, 2, 4), dims=2, num_classes=339, # Hindi writer styles use_checkpoint=True, num_heads=8, num_head_channels=-1, use_scale_shift_norm=True, resblock_updown=False, use_spatial_transformer=True, transformer_depth=1, context_dim=768, text_encoder=self.text_encoder, args=Args() ).to(self.device) self.model.eval() # Count parameters total_params = sum(p.numel() for p in self.model.parameters()) print(f" ✓ UNet initialized ({total_params/1e6:.1f}M parameters)") # Load checkpoint if available if self.checkpoint_path and Path(self.checkpoint_path).exists(): self._load_checkpoint() else: print(f"\n⚠️ No checkpoint found at: {self.checkpoint_path}") print(" Using random initialization") print(f"\n{'='*60}") print(f"✅ Ready for inference!") print(f"{'='*60}\n") except Exception as e: print(f"\n❌ Error during initialization: {str(e)}") raise def _load_checkpoint(self): """Load model checkpoint""" try: print(f"\n📂 Loading checkpoint: {self.checkpoint_path}") checkpoint = torch.load(self.checkpoint_path, map_location=self.device) # Handle different checkpoint formats if isinstance(checkpoint, dict): if 'model_state_dict' in checkpoint: state_dict = checkpoint['model_state_dict'] print(f" Format: Standard (model_state_dict)") elif 'state_dict' in checkpoint: state_dict = checkpoint['state_dict'] print(f" Format: Alternative (state_dict)") else: state_dict = checkpoint print(f" Format: Raw state dict") else: state_dict = checkpoint print(f" Format: Direct tensor state") # Load state dict with strict=False to handle minor mismatches missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False) if missing_keys: print(f" ⚠️ Missing keys: {len(missing_keys)}") if unexpected_keys: print(f" ⚠️ Unexpected keys: {len(unexpected_keys)}") self.checkpoint_loaded = True print(f" ✓ Checkpoint loaded successfully") except Exception as e: print(f" ❌ Failed to load checkpoint: {str(e)}") self.checkpoint_loaded = False def encode_text(self, text): """Encode Hindi text using CANINE""" try: # CANINE handles character-level encoding natively inputs = self.tokenizer( text, return_tensors='pt', padding=True, truncation=True, max_length=512 ) inputs = {k: v.to(self.device) for k, v in inputs.items()} return inputs except Exception as e: print(f"❌ Text encoding error: {e}") return None @torch.no_grad() def generate(self, text, writer_id=0, num_steps=50, guidance_scale=7.5): """ Generate Hindi handwriting from text Args: text: Hindi text in Devanagari script writer_id: Writer style ID (0-338) num_steps: Number of diffusion steps (10-100) guidance_scale: Text guidance strength (1.0-15.0) Returns: Tuple[PIL.Image, str]: Generated image and status message """ if self.model is None: return None, "❌ Model not initialized" try: # Input validation if not text.strip(): return None, "⚠️ Please enter Hindi text" writer_id = max(0, min(int(writer_id), 338)) num_steps = max(10, min(int(num_steps), 100)) guidance_scale = max(1.0, min(float(guidance_scale), 15.0)) print(f"\n🎨 Generating handwriting...") print(f" Text: '{text}'") print(f" Writer: {writer_id}/338") print(f" Steps: {num_steps}") print(f" Guidance: {guidance_scale}") # Encode text with CANINE context = self.encode_text(text) if context is None: return None, "❌ Text encoding failed" batch_size = 1 # Initialize from noise x = torch.randn(batch_size, 1, 64, 64, device=self.device) # Reverse diffusion process for step in range(num_steps - 1, -1, -1): # Prepare timestep and writer conditioning t = torch.full((batch_size,), step, dtype=torch.long, device=self.device) y = torch.tensor([writer_id], dtype=torch.long, device=self.device) # Model prediction with torch.no_grad(): noise_pred = self.model( x, timesteps=t, context=context, y=y ) # Denoising step with adaptive scaling alpha_t = 1.0 - (step / num_steps) scale = guidance_scale * alpha_t x = x - 0.01 * scale * noise_pred # Progress indicator if (num_steps - step) % max(1, num_steps // 5) == 0: progress = ((num_steps - step) / num_steps) * 100 print(f" Progress: {progress:.0f}%") # Post-processing x = torch.clamp(x, -1, 1) x = (x + 1) / 2 # Normalize to [0, 1] x = x.squeeze(0).squeeze(0).cpu().numpy() # Convert to PIL Image img_array = (x * 255).astype(np.uint8) img = Image.fromarray(img_array, mode='L') status = f"✅ Generated with writer {writer_id}, {num_steps} steps" print(f" {status}\n") return img, status except Exception as e: error_msg = f"❌ Generation error: {str(e)}" print(f" {error_msg}") return None, error_msg # ============================================================================== # CONFIGURATION # ============================================================================== # Path to your trained checkpoint (edit this!) CHECKPOINT_PATH = "./checkpoints/model.pt" # Initialize demo print("\n🚀 Initializing DiffusionPen...") demo_instance = DiffusionPenDemo( checkpoint_path=CHECKPOINT_PATH, device=None # Auto-detect GPU/CPU ) def gradio_generate(text, writer_id, num_steps, guidance_scale): """Gradio callback for generation""" img, message = demo_instance.generate( text=text, writer_id=writer_id, num_steps=num_steps, guidance_scale=guidance_scale ) return img, message # ============================================================================== # GRADIO INTERFACE # ============================================================================== theme = gr.themes.Soft( primary_hue="indigo", secondary_hue="amber", ) with gr.Blocks(title="DiffusionPen - Hindi Handwriting Generation", theme=theme) as demo: # Header gr.Markdown(""" # 🎨 DiffusionPen: Hindi Handwriting Generation Generate authentic Hindi handwriting using diffusion models with CANINE text encoding. """) # Main content with gr.Row(): # Input panel with gr.Column(scale=1, min_width=300): gr.Markdown("### ✏️ Input Settings") text_input = gr.Textbox( label="Hindi Text (Devanagari)", placeholder="नमस्ते", lines=2, info="Enter text in Devanagari script" ) writer_id = gr.Slider( label="Writer ID", minimum=0, maximum=338, value=0, step=1, info="0-338: Different writing styles" ) num_steps = gr.Slider( label="Diffusion Steps", minimum=10, maximum=100, value=50, step=10, info="10=fast, 100=quality" ) guidance_scale = gr.Slider( label="Guidance Scale", minimum=1.0, maximum=15.0, value=7.5, step=0.5, info="1=ignore text, 15=strict" ) generate_btn = gr.Button( "✨ Generate Handwriting", variant="primary", size="lg" ) # Output panel with gr.Column(scale=1, min_width=300): gr.Markdown("### 📊 Output") output_image = gr.Image( label="Generated Handwriting", type='pil', interactive=False, show_download_button=True ) status_text = gr.Textbox( label="Status", interactive=False, info="Generation progress and results" ) # Examples gr.Markdown("### 📚 Examples to Try") gr.Examples( examples=[ ["नमस्ते", 0, 50, 7.5], ["हिंदी", 50, 50, 7.5], ["आईआईआीटी", 100, 50, 7.5], ["लिपि", 150, 50, 7.5], ["भाषा", 200, 50, 7.5], ["नई लिखावट", 250, 60, 7.5], ], inputs=[text_input, writer_id, num_steps, guidance_scale], outputs=[output_image, status_text], fn=gradio_generate, cache_examples=False, run_on_click=False ) # Information gr.Markdown(""" --- ### 📖 About This Demo **Model Architecture:** - **Base**: UNet with 128 channels, 3 levels - **Attention**: Spatial transformers at resolutions 16×8 - **Text Encoding**: CANINE (768-dim, character-level) - **Writer Styles**: 339 different writing styles - **Input/Output**: 64×64 grayscale images **CANINE Text Encoder:** - ✓ Character-level (no subword tokenization) - ✓ Native Devanagari support - ✓ Pre-trained on 104 languages - ✓ 768-dimensional contextual embeddings **Performance:** - CPU: ~2 minutes per image - GPU: ~20 seconds per image - Memory: 6-8 GB ### 💡 Tips 1. Keep text short (5-10 characters) for faster generation 2. Try different Writer IDs for style variation 3. Increase steps from 50→100 for better quality 4. Guidance scale 5-10 works best for most cases 5. Use CPU to generate demos, GPU for production ### 🔗 Resources - [CANINE Paper](https://arxiv.org/abs/2103.06367) - [Diffusion Models Course](https://huggingface.co/course) - [UNet Architecture](https://en.wikipedia.org/wiki/U-Net) """) # Connect button generate_btn.click( fn=gradio_generate, inputs=[text_input, writer_id, num_steps, guidance_scale], outputs=[output_image, status_text], api_name="generate" ) if __name__ == "__main__": print(f"\n{'='*60}") print("🚀 Starting DiffusionPen Gradio Demo") print(f"{'='*60}") print(f"Device: {demo_instance.device}") print(f"Checkpoint: {'✓ Loaded' if demo_instance.checkpoint_loaded else '✗ Not found'}") print(f"Models: {'✓ Ready' if demo_instance.model is not None else '✗ Error'}") print(f"{'='*60}\n") demo.launch( share=False, server_name="0.0.0.0", server_port=7860, show_error=True )