Apply Optimized GPU Readback

ADOFAIRenderer.csproj

@@ -5,6 +5,7 @@
         <Nullable>enable</Nullable>
         <NoWarn>MSB3277</NoWarn>
         <EnableSourceControlManagerQueries>false</EnableSourceControlManagerQueries>
+        <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
     </PropertyGroup>
 
     <ItemGroup>

Component/RenderManager.cs

@@ -4,6 +4,7 @@ using System.Collections.Generic;
 using System.IO;
 using System.Threading;
 using ADOFAIRenderer.Util;
+using ADOFAIRenderer.Util.FFmpeg;
 using MelonLoader;
 using MelonLoader.Utils;
 using UnityEngine;
@@ -13,7 +14,7 @@ namespace ADOFAIRenderer.Component
     public class RenderManager : MonoBehaviour
     {
         public static int Countdown;
-        public static int FrameCount;
+        public static long FrameCount;
         public static bool SkipPressToStart;
         public static string LevelPath = "";
         public static RenderManager Instance = null!;
@@ -52,12 +53,12 @@ namespace ADOFAIRenderer.Component
             }
         }
 
-        public static SortedList<int, byte[]> ImageQueue = new SortedList<int, byte[]>();
+        public static SortedList<long, byte[]> ImageQueue = new SortedList<long, byte[]>();
         public static Thread? RenderThread;
-        public static int PassCount;
+        public static long PassCount;
         public static bool EndRequested;
         public static bool EndSuccessful;
-        public static int WonFrame = -1;
+        public static long WonFrame = -1;
 
         private static bool _renderSetup;
 
@@ -125,7 +126,7 @@ namespace ADOFAIRenderer.Component
                     WonFrame = FrameCount;
                 }
 
-                var extraFrames = (int)Math.Ceiling(Configuration.EndWaitTime * Configuration.TargetFPS);
+                var extraFrames = (long)Math.Ceiling(Configuration.EndWaitTime * Configuration.TargetFPS);
                 
                 if (FrameCount - WonFrame < extraFrames) return;
                 RenderMod.IsRendering = RenderMod.IsCapturing = false;
@@ -164,6 +165,8 @@ namespace ADOFAIRenderer.Component
             
             FrameCount = 0;
             
+            RenderMod.FrameReader?.Dispose();
+            RenderMod.FrameReader = new FrameReader();
             StartFFmpegThread();
 
             RenderMod.IsCapturing = true;
@@ -258,7 +261,7 @@ namespace ADOFAIRenderer.Component
                     
                     videoFFmpeg.Start(_currentOutPath, bitrate);
 
-                    var num = isMultiple ? (int)Math.Round(quotient) : 1;
+                    var num = isMultiple ? (long)Math.Round(quotient) : 1L;
 
                     while (true)
                     {
@@ -267,7 +270,7 @@ namespace ADOFAIRenderer.Component
                         {
                             if (ImageQueue.Count > 0)
                             {
-                                int firstKey = ImageQueue.Keys[0];
+                                var firstKey = ImageQueue.Keys[0];
                                 if (firstKey == PassCount)
                                 {
                                     buffer = ImageQueue[firstKey];

Patch/RenderEnvironmentPatches.cs

@@ -87,7 +87,7 @@ namespace ADOFAIRenderer.Patch
             public static bool Prefix(ref int __result)
             {
                 if (!RenderMod.IsCapturing || !RenderPatch) return true;
-                __result = RenderManager.FrameCount + RenderMod.FrameCountBase;
+                __result = (int)(RenderManager.FrameCount + RenderMod.FrameCountBase);
                 return false;
             }
         }

RenderMod.cs

@@ -3,6 +3,7 @@ using System.Reflection;
 using ADOFAIRenderer;
 using ADOFAIRenderer.Component;
 using ADOFAIRenderer.Util;
+using ADOFAIRenderer.Util.FFmpeg;
 using MelonLoader;
 using UnityEngine;
 using UnityEngine.EventSystems;
@@ -38,6 +39,8 @@ namespace ADOFAIRenderer
         public static double DspTimeBase;
         public static int FrameCountBase;
 
+        public static FrameReader? FrameReader;
+
         public static bool RenderCheck = true;
         public static int SkipCount;
 
@@ -98,7 +101,7 @@ namespace ADOFAIRenderer
                 var targetFps = RenderManager.Configuration.TargetFPS;
                 var videoFps = RenderManager.Configuration.VideoFPS;
                 var quotient = targetFps / videoFps;
-                var num = Math.Abs(quotient - Math.Round(quotient)) < 1e-9 ? (int)Math.Round(quotient) : 1;
+                var num = Math.Abs(quotient - Math.Round(quotient)) < 1e-9 ? (long)Math.Round(quotient) : 1L;
 
                 if (FrameCountForVideo % num == 0)
                 {
@@ -126,7 +129,7 @@ namespace ADOFAIRenderer
             }
         }
 
-        public static int FrameCountForVideo;
+        public static long FrameCountForVideo;
         public static int ReadbackFailureCount;
         public static int ThrottleCount;
         public static bool IsThrottling;
@@ -172,6 +175,9 @@ namespace ADOFAIRenderer
                 return;
             }
             
+            FrameReader?.RequestFrameCapture(texture, currentFrame);
+
+            /*
             AsyncGPUReadback.Request(texture, 0, delegate(AsyncGPUReadbackRequest req)
             {
                 if (!IsCapturing) return;
@@ -202,7 +208,7 @@ namespace ADOFAIRenderer
                 {
                     RenderManager.ImageQueue[currentFrame] = rawBytes;
                 }
-            });
+            });*/
         }
 
         public static void StartRender()

UI/RenderButton.cs

@@ -1,5 +1,6 @@
 using System.Linq;
 using ADOFAIRenderer.Component;
+using ADOFAIRenderer.Util.FFmpeg;
 using MelonLoader;
 using TMPro;
 using UnityEngine;
@@ -35,7 +36,7 @@ namespace ADOFAIRenderer.UI
             if (codecDropdown == null) return;
             codecDropdown.ClearOptions();
             
-            var options = System.Enum.GetNames(typeof(Util.VideoFFmpeg.Codec)).ToList();
+            var options = System.Enum.GetNames(typeof(VideoFFmpeg.Codec)).ToList();
             codecDropdown.AddOptions(options);
         }
 

Util/FFmpeg/FFmpegProcess.cs

@@ -52,15 +52,21 @@ namespace ADOFAIRenderer.Util.FFmpeg
             _ffmpegProcess.WaitForExit();
         }
 
-        public static byte[] RentRawBytes<T>(NativeArray<T> arr) where T : struct
+        public static byte[] RentRawBytes(int size)
+        {
+            var array = ArrayPool<byte>.Shared.Rent(size);
+            return array;
+        }
+
+        /*public static byte[] RentRawBytes<T>(NativeArray<T> arr) where T : struct
         {
             var val = new NativeSlice<T>(arr).SliceConvert<byte>();
             var array = ArrayPool<byte>.Shared.Rent(val.Length);
             val.CopyTo(array);
             return array;
-        }
+        }*/
 
-        public static byte[] ToRawBytes<T>(NativeArray<T> arr) where T : struct
+        protected static byte[] ToRawBytes<T>(NativeArray<T> arr) where T : struct
         {
             var val = new NativeSlice<T>(arr).SliceConvert<byte>();
             var array = new byte[val.Length];

Util/FFmpeg/VideoFFmpeg.cs

@@ -1,7 +1,6 @@
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using Unity.Collections;
-using UnityEngine;
 
 namespace ADOFAIRenderer.Util.FFmpeg
 {

Util/FrameReader.cs

@@ -1,7 +1,181 @@
-namespace ADOFAIRenderer.Util
+using System;
+using System.Collections.Generic;
+using ADOFAIRenderer.Component;
+using ADOFAIRenderer.Util.FFmpeg;
+using MelonLoader;
+using Unity.Collections;
+using Unity.Collections.LowLevel.Unsafe;
+using UnityEngine;
+using UnityEngine.Rendering;
+
+namespace ADOFAIRenderer.Util
 {
-    public class FrameReader
+    public class FrameReader : IDisposable
+    {
+        private const int InitialBufferSize = 16;
+        
+        private readonly Stack<ReadbackRequestState> _statePool = new Stack<ReadbackRequestState>();
+        private readonly object _poolLock = new object();
+        private bool _isInitialized;
+        
+        public FrameReader()
+        {
+            InitializePool();
+        }
+
+        private void InitializePool()
         {
             
+            lock (_poolLock)
+            {
+                if (_isInitialized) return;
+                for (var i = 0; i < InitialBufferSize; i++)
+                {
+                    _statePool.Push(new ReadbackRequestState(this));
+                }
+
+                _isInitialized = true;
+            }
+        }
+
+        public void RequestFrameCapture(Texture texture, long currentFrame)
+        {
+            if (!RenderMod.IsCapturing) return;
+            ReadbackRequestState state;
+
+            lock (_poolLock)
+            {
+                if (_statePool.Count == 0)
+                {
+                    state = new ReadbackRequestState(this);
+                }
+                else
+                {
+                    state = _statePool.Pop();
+                }
+            }
+
+            var bufferSize = texture.width * texture.height * 4; // RGBA = 4 channel
+            state.Setup(currentFrame, bufferSize);
+
+            AsyncGPUReadback.RequestIntoNativeArray(
+                ref state.NativeBuffer,
+                texture,
+                0,
+                state.CachedCallback
+            );
+        }
+
+        internal void OnReadbackComplete(ReadbackRequestState state, AsyncGPUReadbackRequest request)
+        {
+            var frame = state.FrameIndex;
+
+            if (!RenderMod.IsCapturing)
+            {
+                RecycleState(state);
+                return;
+            }
+
+            if (request.hasError)
+            {
+                RenderMod.ReadbackFailureCount++;
+                Melon<RenderMod>.Logger.Warning($"[OOM Prevent] GPU Readback failed for frame {frame}. Queue deadlock prevented by sending an empty frame.");
+
+                PushToQueue(frame, Array.Empty<byte>());
+            }
+            else
+            {
+                unsafe
+                {
+                    var ptr = state.NativeBuffer.GetUnsafeReadOnlyPtr();
+                    var rawBytes = FFmpegProcess.RentRawBytes(state.NativeBuffer.Length);
+
+                    fixed (byte* destPtr = rawBytes)
+                    {
+                        Buffer.MemoryCopy(ptr, destPtr, rawBytes.Length, state.NativeBuffer.Length);
+                    }
+
+                    PushToQueue(frame, rawBytes);
+                }
+            }
+
+            RecycleState(state);
+        }
+
+        private void RecycleState(ReadbackRequestState state)
+        {
+            state.Recycle();
+            lock (_poolLock)
+            {
+                _statePool.Push(state);
+            }
+        }
+        
+        private static void PushToQueue(long frame, byte[] rawBytes)
+        {
+            lock (RenderManager.ImageQueue)
+            {
+                RenderManager.ImageQueue[frame] = rawBytes;
+            }
+        }
+        
+        public void Dispose()
+        {
+            lock (_poolLock)
+            {
+                while (_statePool.Count > 0)
+                {
+                    var state = _statePool.Pop();
+                    state.Dispose();
+                }
+                _isInitialized = false;
+            }
+        }
+    }
+
+    internal class ReadbackRequestState : IDisposable
+    {
+        private readonly FrameReader _owner;
+        public long FrameIndex { get; private set; }
+        public NativeArray<byte> NativeBuffer;
+
+        public readonly Action<AsyncGPUReadbackRequest> CachedCallback;
+
+        public ReadbackRequestState(FrameReader owner)
+        {
+            _owner = owner;
+            CachedCallback = OnComplete;
+            FrameIndex = -1;
+        }
+
+        public void Setup(long frameIndex, int bufferSize)
+        {
+            FrameIndex = frameIndex;
+            
+            // 버퍼가 생성되지 않았거나 크기가 다른 경우 재생성
+            if (!NativeBuffer.IsCreated || NativeBuffer.Length != bufferSize)
+            {
+                // 이전 메모리 해제
+                if (NativeBuffer.IsCreated) NativeBuffer.Dispose();
+
+                NativeBuffer = new NativeArray<byte>(bufferSize, Allocator.Persistent,
+                    NativeArrayOptions.UninitializedMemory);
+            }
+        }
+
+        private void OnComplete(AsyncGPUReadbackRequest request)
+        {
+            _owner.OnReadbackComplete(this, request);
+        }
+
+        public void Recycle()
+        {
+            FrameIndex = -1;
+        }
+
+        public void Dispose()
+        {
+            if (NativeBuffer.IsCreated) NativeBuffer.Dispose();
+        }
     }
 }