File size: 2,227 Bytes
5bcb990 b2c5af8 5bcb990 b2c5af8 5bcb990 b2c5af8 5bcb990 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import torch
from PIL import Image
import torchvision.transforms as transforms
import logging
from datetime import datetime
import io
logger = logging.getLogger(__name__)
class AnimeGANProcessor:
def __init__(self, device):
self.device = device
self.model = None
self.load_model()
def load_model(self):
try:
logger.info("Loading AnimeGAN model...")
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:
logger.error(f"Error loading model: {str(e)}")
raise
def process_image(self, image):
"""Convert image to anime style"""
try:
transform = transforms.Compose([
transforms.Resize((512, 512)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
with torch.no_grad():
output = self.model(transform(image).unsqueeze(0).to(self.device))
return transforms.ToPILImage()((output * 0.5 + 0.5).squeeze().cpu())
except Exception as e:
logger.error(f"Error processing image: {str(e)}")
raise
def generate_anime(image_data):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
processor = AnimeGANProcessor(device)
start_time = datetime.now()
logger.info(f"Generating anime image - {start_time}")
try:
image = Image.open(io.BytesIO(image_data)).convert("RGB")
processed_img = processor.process_image(image)
img_io = io.BytesIO()
processed_img.save(img_io, 'PNG')
img_io.seek(0)
duration = (datetime.now() - start_time).total_seconds()
logger.info(f"Successfully processed. Duration: {duration} seconds")
return img_io
except Exception as e:
logger.error(f"Processing error: {str(e)}", exc_info=True)
raise |