diff --git a/Plugins/StyleTransfer/Content/T_StyleImage.uasset b/Plugins/StyleTransfer/Content/T_StyleImage.uasset index 33212ee9..0cb846b9 100644 --- a/Plugins/StyleTransfer/Content/T_StyleImage.uasset +++ b/Plugins/StyleTransfer/Content/T_StyleImage.uasset @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:936df5dec6fff94f1da76ce02bac39fcaa46848fe7a760c7f9c2fea6b4cdcfb8 -size 2714591 +oid sha256:acbfb62d054055fd0d80e5f664a1525994296d41dfe3da89527deac01f2341bd +size 2714563 diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp index 07024ee7..9e4033d0 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp @@ -113,6 +113,7 @@ void FStyleTransferSceneViewExtension::SubscribeToPostProcessingPass(EPostProces void FStyleTransferSceneViewExtension::PreRenderViewFamily_RenderThread(FRDGBuilder& GraphBuilder, FSceneViewFamily& InViewFamily) { + return; const FName RenderCaptureProviderType = IRenderCaptureProvider::GetModularFeatureName(); if(!IModularFeatures::Get().IsModularFeatureAvailable(RenderCaptureProviderType)) return; @@ -181,19 +182,19 @@ FRDGTexture* FStyleTransferSceneViewExtension::TensorToTexture(FRDGBuilder& Grap return OutputTexture; } -void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder, const FScreenPassTexture& SourceTexture, const FNeuralTensor& DestinationTensor) +void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor) { const FIntVector InputTensorDimensions = { CastNarrowingSafe(DestinationTensor.GetSize(1)), CastNarrowingSafe(DestinationTensor.GetSize(2)), CastNarrowingSafe(DestinationTensor.GetSize(3)), }; - const FIntPoint SceneColorRenderTargetDimensions = SourceTexture.Texture->Desc.Extent; + const FIntPoint SceneColorRenderTargetDimensions = SourceTexture->Desc.Extent; FRDGBufferRef StyleTransferContentInputBuffer = GraphBuilder.RegisterExternalBuffer(DestinationTensor.GetPooledBuffer()); auto SceneColorToInputTensorParameters = GraphBuilder.AllocParameters(); SceneColorToInputTensorParameters->TensorVolume = CastNarrowingSafe(DestinationTensor.Num()); - SceneColorToInputTensorParameters->InputTexture = SourceTexture.Texture; + SceneColorToInputTensorParameters->InputTexture = SourceTexture; SceneColorToInputTensorParameters->InputTextureSampler = TStaticSamplerState::GetRHI(); SceneColorToInputTensorParameters->OutputUAV = GraphBuilder.CreateUAV(StyleTransferContentInputBuffer); SceneColorToInputTensorParameters->OutputDimensions = {InputTensorDimensions.X, InputTensorDimensions.Y}; @@ -239,7 +240,7 @@ FScreenPassTexture FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap const FNeuralTensor& StyleTransferContentInputTensor = StyleTransferNetwork->GetInputTensorForContext(*InferenceContext, 0); - TextureToTensor(GraphBuilder, SceneColor, StyleTransferContentInputTensor); + TextureToTensor(GraphBuilder, SceneColor.Texture, StyleTransferContentInputTensor); StyleTransferNetwork->Run(GraphBuilder, *InferenceContext); diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp index a21aea1e..8d947493 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp @@ -3,11 +3,13 @@ #include "StyleTransferSubsystem.h" +#include "IRenderCaptureProvider.h" #include "NeuralNetwork.h" #include "RenderGraphUtils.h" #include "ScreenPass.h" #include "StyleTransferModule.h" #include "StyleTransferSceneViewExtension.h" +#include "TextureCompiler.h" #include "Rendering/Texture2DResource.h" TAutoConsoleVariable CVarStyleTransferEnabled( @@ -47,11 +49,14 @@ void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportCl checkf(*StyleTransferInferenceContext != INDEX_NONE, TEXT("Could not create inference context for StyleTransferNetwork")); UTexture2D* StyleTexture = LoadObject(this, TEXT("/Script/Engine.Texture2D'/StyleTransfer/T_StyleImage.T_StyleImage'")); - UpdateStyle(StyleTexture); + FTextureCompilingManager::Get().FinishCompilation({StyleTexture}); + //UpdateStyle(StyleTexture); + UpdateStyle(FPaths::GetPath("C:\\projects\\realtime-style-transfer\\temp\\style_params_tensor.bin")); StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension(ViewportClient, StyleTransferNetwork, StyleTransferInferenceContext.ToSharedRef()); } StyleTransferSceneViewExtension->SetEnabled(true); + ViewportClient->GetWorld()->GetWorldSettings()->SetPauserPlayerState(ViewportClient->GetWorld()->GetFirstPlayerController()->PlayerState); } void UStyleTransferSubsystem::StopStylizingViewport() @@ -74,43 +79,75 @@ void UStyleTransferSubsystem::UpdateStyle(UTexture2D* StyleTexture) { checkf(StyleTransferInferenceContext.IsValid() && (*StyleTransferInferenceContext) != INDEX_NONE, TEXT("Can not infer style without inference context")); checkf(StylePredictionInferenceContext != INDEX_NONE, TEXT("Can not update style without inference context")); - + FlushRenderingCommands(); ENQUEUE_RENDER_COMMAND(StylePrediction)([this, StyleTexture](FRHICommandListImmediate& RHICommandList) { - FRDGBuilder GraphBuilder(RHICommandList); - - - const FNeuralTensor& InputStyleImageTensor = StylePredictionNetwork->GetInputTensorForContext(StylePredictionInferenceContext, 0); - - FTextureResource* StyleTextureResource = StyleTexture->GetResource(); - if(!StyleTextureResource->IsInitialized()) + IRenderCaptureProvider* RenderCaptureProvider = nullptr; + const FName RenderCaptureProviderType = IRenderCaptureProvider::GetModularFeatureName(); + if(IModularFeatures::Get().IsModularFeatureAvailable(RenderCaptureProviderType)) { - StyleTextureResource->UpdateRHI(); + RenderCaptureProvider = &IModularFeatures::Get().GetModularFeature(RenderCaptureProviderType); + RenderCaptureProvider->BeginCapture(&RHICommandList); } - FRDGTextureRef RDGStyleTexture = GraphBuilder.RegisterExternalTexture(CreateRenderTarget(StyleTextureResource->TextureRHI, TEXT("StyleInputTexture"))); - FStyleTransferSceneViewExtension::TextureToTensor(GraphBuilder, FScreenPassTexture(RDGStyleTexture), InputStyleImageTensor); + FRDGBuilder GraphBuilder(RHICommandList); + { + RDG_EVENT_SCOPE(GraphBuilder, "StylePrediction"); - StylePredictionNetwork->Run(GraphBuilder, StylePredictionInferenceContext); + const FNeuralTensor& InputStyleImageTensor = StylePredictionNetwork->GetInputTensorForContext(StylePredictionInferenceContext, 0); + FTextureResource* StyleTextureResource = StyleTexture->GetResource(); + FRDGTextureRef RDGStyleTexture = GraphBuilder.RegisterExternalTexture(CreateRenderTarget(StyleTextureResource->TextureRHI, TEXT("StyleInputTexture"))); + FStyleTransferSceneViewExtension::TextureToTensor(GraphBuilder, RDGStyleTexture, InputStyleImageTensor); - const FNeuralTensor& OutputStyleParams = StylePredictionNetwork->GetOutputTensorForContext(StylePredictionInferenceContext, 0); - const FNeuralTensor& InputStyleParams = StyleTransferNetwork->GetInputTensorForContext(*StyleTransferInferenceContext, StyleTransferStyleParamsInputIndex); + StylePredictionNetwork->Run(GraphBuilder, StylePredictionInferenceContext); - FRDGBufferRef OutputStyleParamsBuffer = GraphBuilder.RegisterExternalBuffer(OutputStyleParams.GetPooledBuffer()); - FRDGBufferRef InputStyleParamsBuffer = GraphBuilder.RegisterExternalBuffer(InputStyleParams.GetPooledBuffer()); - const uint64 NumBytes = OutputStyleParams.NumInBytes(); - check(OutputStyleParamsBuffer->GetSize() == InputStyleParamsBuffer->GetSize()); - check(OutputStyleParamsBuffer->GetSize() == OutputStyleParams.NumInBytes()); - check(InputStyleParamsBuffer->GetSize() == InputStyleParams.NumInBytes()); + const FNeuralTensor& OutputStyleParams = StylePredictionNetwork->GetOutputTensorForContext(StylePredictionInferenceContext, 0); + const FNeuralTensor& InputStyleParams = StyleTransferNetwork->GetInputTensorForContext(*StyleTransferInferenceContext, StyleTransferStyleParamsInputIndex); - AddCopyBufferPass(GraphBuilder, InputStyleParamsBuffer, OutputStyleParamsBuffer); + FRDGBufferRef OutputStyleParamsBuffer = GraphBuilder.RegisterExternalBuffer(OutputStyleParams.GetPooledBuffer()); + FRDGBufferRef InputStyleParamsBuffer = GraphBuilder.RegisterExternalBuffer(InputStyleParams.GetPooledBuffer()); + const uint64 NumBytes = OutputStyleParams.NumInBytes(); + check(OutputStyleParamsBuffer->GetSize() == InputStyleParamsBuffer->GetSize()); + check(OutputStyleParamsBuffer->GetSize() == OutputStyleParams.NumInBytes()); + check(InputStyleParamsBuffer->GetSize() == InputStyleParams.NumInBytes()); + AddCopyBufferPass(GraphBuilder, InputStyleParamsBuffer, OutputStyleParamsBuffer); + } GraphBuilder.Execute(); + + if(RenderCaptureProvider) + { + RenderCaptureProvider->EndCapture(&RHICommandList); + } }); FlushRenderingCommands(); } +void UStyleTransferSubsystem::UpdateStyle(FString StyleTensorDataPath) +{ + FArchive& FileReader = *IFileManager::Get().CreateFileReader(*StyleTensorDataPath); + TArray StyleParams; + StyleParams.SetNumUninitialized(192); + + FileReader << StyleParams; + + ENQUEUE_RENDER_COMMAND(StyleParamsLoad)([this, StyleParams = MoveTemp(StyleParams)](FRHICommandListImmediate& RHICommandList) + { + const FNeuralTensor& InputStyleParams = StyleTransferNetwork->GetInputTensorForContext(*StyleTransferInferenceContext, StyleTransferStyleParamsInputIndex); + + FRDGBuilder GraphBuilder(RHICommandList); + { + RDG_EVENT_SCOPE(GraphBuilder, "StyleParamsLoad"); + + FRDGBufferRef InputStyleParamsBuffer = GraphBuilder.RegisterExternalBuffer(InputStyleParams.GetPooledBuffer()); + GraphBuilder.QueueBufferUpload(InputStyleParamsBuffer, StyleParams.GetData(), StyleParams.Num() * StyleParams.GetTypeSize(), ERDGInitialDataFlags::NoCopy); + } + GraphBuilder.Execute(); + }); + FlushRenderingCommands(); +} + void UStyleTransferSubsystem::HandleConsoleVariableChanged(IConsoleVariable* ConsoleVariable) { check(ConsoleVariable == CVarStyleTransferEnabled.AsVariable()); diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h index 7f161c8e..5f05b458 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h @@ -38,7 +38,7 @@ public: static void AddRescalingTextureCopy(FRDGBuilder& GraphBuilder, FRDGTexture& RDGSourceTexture, FScreenPassRenderTarget& DestinationRenderTarget); static FRDGTexture* TensorToTexture(FRDGBuilder& GraphBuilder, const FRDGTextureDesc& BaseDestinationDesc, const FNeuralTensor& SourceTensor); - static void TextureToTensor(FRDGBuilder& GraphBuilder, const FScreenPassTexture& SourceTexture, const FNeuralTensor& DestinationTensor); + static void TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor); private: /** The actual Network pointer is not tracked so we need a WeakPtr too so we can check its validity on the game thread. */ diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h index bf419bd5..098a06f8 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h @@ -26,6 +26,7 @@ public: void StopStylizingViewport(); void UpdateStyle(UTexture2D* StyleTexture); + void UpdateStyle(FString StyleTensorDataPath); private: FStyleTransferSceneViewExtension::Ptr StyleTransferSceneViewExtension;