Create inference.py

FantasyTalking/inference.py

@@ -0,0 +1,50 @@
+import os
+import torch
+from PIL import Image
+import torchvision.transforms as transforms
+from fantasy_talking.model import FantasyTalkingModel
+from moviepy.editor import ImageSequenceClip
+import torchaudio
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# تحميل الموديل
+model_ckpt = "./models/fantasytalking_model.ckpt"
+model = FantasyTalkingModel()
+model.load_state_dict(torch.load(model_ckpt, map_location=device))
+model = model.to(device)
+model.eval()
+
+# تحويل الصورة إلى Tensor
+def load_image(image_path):
+    image = Image.open(image_path).convert("RGB")
+    transform = transforms.Compose([
+        transforms.Resize((512, 512)),
+        transforms.ToTensor()
+    ])
+    return transform(image).unsqueeze(0).to(device)
+
+# تحويل الصوت إلى Tensor
+def load_audio(audio_path):
+    waveform, sample_rate = torchaudio.load(audio_path)
+    if waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+    if sample_rate != 16000:
+        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
+        waveform = resampler(waveform)
+    return waveform.to(device), 16000
+
+# توليد الفيديو
+def generate_video(image_path, audio_path, output_path="output.mp4"):
+    image_tensor = load_image(image_path)
+    audio_tensor, _ = load_audio(audio_path)
+
+    with torch.no_grad():
+        frames = model.generate(image_tensor, audio_tensor)
+
+    # حفظ الفيديو من الفريمات
+    frames = [transforms.ToPILImage()(frame.squeeze(0).cpu()) for frame in frames]
+    video_clip = ImageSequenceClip([frame for frame in frames], fps=25)
+    video_clip.write_videofile(output_path, codec="libx264")
+
+    return output_path