Update demo_generate/animegan_method.py

demo_generate/animegan_method.py

@@ -3,6 +3,7 @@ from PIL import Image
 import torchvision.transforms as transforms
 import logging
 from datetime import datetime
+import io
 
 logger = logging.getLogger(__name__)
 
@@ -15,8 +16,12 @@ class AnimeGANProcessor:
     def load_model(self):
         try:
             logger.info("Loading AnimeGAN model...")
-            self.model = torch.hub.load('bryandlee/animegan2-pytorch:main', 'generator', trust_repo=True).to(self.device)
-            self.model.load_state_dict(torch.load('face_paint_512_v2.pt', map_location=self.device))
+            self.model = torch.hub.load(
+                'bryandlee/animegan2-pytorch:main', 
+                'generator',
+                pretrained='face_paint_512_v2',
+                trust_repo=True
+            ).to(self.device)
             self.model.eval()
             logger.info("Model loaded successfully")
         except Exception as e:
@@ -24,6 +29,7 @@ class AnimeGANProcessor:
             raise
     
     def process_image(self, image):
+        """Convert image to anime style"""
         try:
             transform = transforms.Compose([
                 transforms.Resize((512, 512)),