Create inference.py

inference.py

@@ -0,0 +1,125 @@
+import argparse
+import os
+import torch
+from data_loader.loader import ContentData
+from models.unet import UNetModel
+from diffusers import AutoencoderKL
+from models.diffusion import Diffusion
+import torchvision
+from parse_config import cfg, cfg_from_file, assert_and_infer_cfg
+from utils.util import fix_seed
+from PIL import Image
+import torchvision.transforms as transforms
+
+class OneDMInference:
+    def __init__(self, model_path, cfg_path, device='cuda' if torch.cuda.is_available() else 'cpu'):
+        self.device = device
+        
+        # Load config
+        cfg_from_file(cfg_path)
+        assert_and_infer_cfg()
+        fix_seed(cfg.TRAIN.SEED)
+        
+        # Initialize models
+        self.unet = self._initialize_unet(model_path)
+        self.vae = self._initialize_vae()
+        self.diffusion = Diffusion(device=self.device)
+        self.content_loader = ContentData()
+        
+        # Define transform
+        self.transform = transforms.Compose([
+            transforms.Grayscale(),
+            transforms.ToTensor()
+        ])
+    
+    def _initialize_unet(self, model_path):
+        unet = UNetModel(
+            in_channels=cfg.MODEL.IN_CHANNELS,
+            model_channels=cfg.MODEL.EMB_DIM,
+            out_channels=cfg.MODEL.OUT_CHANNELS,
+            num_res_blocks=cfg.MODEL.NUM_RES_BLOCKS,
+            attention_resolutions=(1,1),
+            channel_mult=(1, 1),
+            num_heads=cfg.MODEL.NUM_HEADS,
+            context_dim=cfg.MODEL.EMB_DIM
+        ).to(self.device)
+        
+        # Load model with weights_only=True
+        unet.load_state_dict(torch.load(model_path, map_location=self.device, weights_only=True))
+        unet.eval()
+        return unet
+    
+    def _initialize_vae(self):
+        vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
+        vae = vae.to(self.device)
+        vae.requires_grad_(False)
+        return vae
+
+    def _load_image(self, image_path):
+        image = Image.open(image_path)
+        image_tensor = self.transform(image)
+        return image_tensor
+    
+    def generate(self, text, style_path, laplace_path, output_dir, 
+                sample_method='ddim', sampling_timesteps=50, eta=0.0):
+        """
+        Generate handwritten text with the specified style
+        """
+        # Load style and laplace images
+        style_input = self._load_image(style_path).unsqueeze(0).to(self.device)
+        laplace = self._load_image(laplace_path).unsqueeze(0).to(self.device)
+        
+        # Prepare text reference
+        text_ref = self.content_loader.get_content(text)
+        text_ref = text_ref.to(self.device).repeat(1, 1, 1, 1)
+        
+        # Initialize noise
+        x = torch.randn((text_ref.shape[0], 4, style_input.shape[2]//8, 
+                        (text_ref.shape[1]*32)//8)).to(self.device)
+        
+        # Generate image
+        if sample_method == 'ddim':
+            sampled_images = self.diffusion.ddim_sample(
+                self.unet, self.vae, style_input.shape[0],
+                x, style_input, laplace, text_ref,
+                sampling_timesteps, eta
+            )
+        elif sample_method == 'ddpm':
+            sampled_images = self.diffusion.ddpm_sample(
+                self.unet, self.vae, style_input.shape[0],
+                x, style_input, laplace, text_ref
+            )
+        
+        # Save generated image
+        os.makedirs(output_dir, exist_ok=True)
+        output_paths = []
+        for idx, image in enumerate(sampled_images):
+            im = torchvision.transforms.ToPILImage()(image)
+            image = im.convert("L")
+            output_path = os.path.join(output_dir, f"{text}_{idx}.png")
+            image.save(output_path)
+            output_paths.append(output_path)
+            
+        return output_paths
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_path', required=True, help='Path to the One-DM model checkpoint')
+    parser.add_argument('--cfg_path', required=True, help='Path to the config file')
+    parser.add_argument('--text', required=True, help='Text to generate')
+    parser.add_argument('--style_path', required=True, help='Path to style image')
+    parser.add_argument('--laplace_path', required=True, help='Path to laplace image')
+    parser.add_argument('--output_dir', required=True, help='Output directory')
+    args = parser.parse_args()
+    
+    model = OneDMInference(args.model_path, args.cfg_path)
+    output_paths = model.generate(
+        args.text,
+        args.style_path,
+        args.laplace_path,
+        args.output_dir
+    )
+    print(f"Generated images saved at: {output_paths}")
+
+if __name__ == "__main__":
+    main()