diff --git a/Plugins/GameFeatures/ShooterCore/Content/Experiences/B_ShooterGame_Elimination.uasset b/Plugins/GameFeatures/ShooterCore/Content/Experiences/B_ShooterGame_Elimination.uasset index 80c7c83b..5419244c 100644 --- a/Plugins/GameFeatures/ShooterCore/Content/Experiences/B_ShooterGame_Elimination.uasset +++ b/Plugins/GameFeatures/ShooterCore/Content/Experiences/B_ShooterGame_Elimination.uasset @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:d560527d0c8d9b3e638d06665c5f1376d8239441c83268fa35c97cf63ed9f497 -size 11675 +oid sha256:8e6585c67bdb4a197e6783b3da6fc738f9b2b5a514baf16b43b6f1ed512c576e +size 11475 diff --git a/Plugins/StyleTransfer/Content/NN_StylePredictor.uasset b/Plugins/StyleTransfer/Content/NN_StylePredictor.uasset index 65f999e0..668c7c07 100644 --- a/Plugins/StyleTransfer/Content/NN_StylePredictor.uasset +++ b/Plugins/StyleTransfer/Content/NN_StylePredictor.uasset @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:5ffdd6afea702024f6fa2d00f2dbf90bef02fc307e61c9382b0cbeba21282059 -size 4092053 +oid sha256:e4ed7196066ce72fc937861610095380a4d89be933b99f82db22ec1d96696347 +size 4092063 diff --git a/Plugins/StyleTransfer/Content/NN_StyleTransfer.uasset b/Plugins/StyleTransfer/Content/NN_StyleTransfer.uasset index 7b862770..0556bbbe 100644 --- a/Plugins/StyleTransfer/Content/NN_StyleTransfer.uasset +++ b/Plugins/StyleTransfer/Content/NN_StyleTransfer.uasset @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:7c80d8e6510bae82db869b7b0175b20ee71163b8de7dec675eead1c282944129 -size 6700104 +oid sha256:1682560d011ef8c769ea03bc18f2b95849edb2e0e7deab9b62f786ddfbf2d2bc +size 6699914 diff --git a/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf b/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf index 47fa3a25..a41de6f2 100644 --- a/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf +++ b/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf @@ -3,23 +3,21 @@ #include "/Engine/Public/Platform.ush" +// this assumes that the OutputTexture has +// the exact same dimensions as InputTensor! +uint TensorVolume; +uint2 TextureSize; RWTexture2D OutputTexture; Buffer InputTensor; +// DispatchThreadID corresponds to InputTensor shape dimensions not texture XY -> DispatchThreadID.X = Texture.Y [numthreads(THREADGROUP_SIZE_X, THREADGROUP_SIZE_Y, THREADGROUP_SIZE_Z)] void OutputTensorToSceneColorCS(in const uint3 DispatchThreadID : SV_DispatchThreadID) { - - uint TensorVolume; - InputTensor.GetDimensions(TensorVolume); - - uint2 TextureSize = 0; - OutputTexture.GetDimensions(TextureSize.x, TextureSize.y); - // note that the input tensor has shape (1, Y, X, C) // which is why we need to flip the indexing - const uint PixelIndex = DispatchThreadID.x * TextureSize.y + DispatchThreadID.x; - const uint GlobalIndex = PixelIndex * 3; + const uint TensorPixelNumber = DispatchThreadID.x * TextureSize.x + DispatchThreadID.y; + const uint GlobalIndex = TensorPixelNumber * 3; if (GlobalIndex >= TensorVolume) { @@ -27,10 +25,11 @@ void OutputTensorToSceneColorCS(in const uint3 DispatchThreadID : SV_DispatchThr } const uint2 TextureCoords = uint2(DispatchThreadID.y, DispatchThreadID.x); - OutputTexture[TextureCoords] = float4( + const float4 RGBAColor = float4( InputTensor[GlobalIndex + 0], InputTensor[GlobalIndex + 1], InputTensor[GlobalIndex + 2], - 1.0f + 0.0f ); + OutputTexture[TextureCoords] = RGBAColor; } diff --git a/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf b/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf index ec794534..f762454a 100644 --- a/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf +++ b/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf @@ -12,16 +12,16 @@ float2 HalfPixelUV; void SceneColorToInputTensorCS(in const uint3 DispatchThreadID : SV_DispatchThreadID) { const uint2 OutputUAVTexelCoordinate = DispatchThreadID.xy; - if(any(OutputUAVTexelCoordinate > OutputDimensions)) + if(any(OutputUAVTexelCoordinate >= OutputDimensions)) { return; } - const uint GlobalIndex = (OutputUAVTexelCoordinate.x * OutputDimensions.x + OutputUAVTexelCoordinate.y) * 3; + const uint GlobalIndex = ((OutputUAVTexelCoordinate.x * OutputDimensions.y + OutputUAVTexelCoordinate.y)*3); - // note that the input tensor has shape (1, Y, X, C) + // note that the OutputUAV has shape (1, Y, X, C) // which is why we need to flip the indexing - const float2 UV = OutputUAVTexelCoordinate.yx / OutputDimensions.yx + HalfPixelUV; + const float2 UV = float2(OutputUAVTexelCoordinate.yx) / float2(OutputDimensions.yx) + HalfPixelUV; const float4 TextureValue = InputTexture.SampleLevel(InputTextureSampler, UV, 0); diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp index 38a6295b..24685448 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp @@ -18,7 +18,9 @@ #include "Containers/DynamicRHIResourceArray.h" #include "PostProcess/PostProcessMaterial.h" #include "OutputTensorToSceneColorCS.h" +#include "PixelShaderUtils.h" #include "SceneColorToInputTensorCS.h" +#include "StyleTransferModule.h" template OutType CastNarrowingSafe(InType InValue) @@ -35,8 +37,9 @@ OutType CastNarrowingSafe(InType InValue) } -FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, int32 InInferenceContext) +FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef InInferenceContext) : FSceneViewExtensionBase(AutoRegister) + , StyleTransferNetworkWeakPtr(InStyleTransferNetwork) , StyleTransferNetwork(InStyleTransferNetwork) , LinkedViewportClient(AssociatedViewportClient) , InferenceContext(InInferenceContext) @@ -48,9 +51,54 @@ void FStyleTransferSceneViewExtension::SetupViewFamily(FSceneViewFamily& InViewF { } -void FStyleTransferSceneViewExtension::SubscribeToPostProcessingPass(EPostProcessingPass PassId, - FAfterPassCallbackDelegateArray& - InOutPassCallbacks, bool bIsPassEnabled) +bool FStyleTransferSceneViewExtension::IsActiveThisFrame_Internal(const FSceneViewExtensionContext& Context) const +{ + check(IsInGameThread()); + return bIsEnabled && *InferenceContext != -1 && StyleTransferNetworkWeakPtr.IsValid(); +} + +void FStyleTransferSceneViewExtension::AddRescalingTextureCopy(FRDGBuilder& GraphBuilder, FRDGTexture& RDGSourceTexture, FScreenPassRenderTarget& DestinationRenderTarget) +{ + FGlobalShaderMap* ShaderMap = GetGlobalShaderMap(GMaxRHIFeatureLevel); + + TShaderMapRef VertexShader(ShaderMap); + + TShaderMapRef PixelShader(ShaderMap); + + FCopyRectPS::FParameters* PixelShaderParameters = GraphBuilder.AllocParameters(); + PixelShaderParameters->InputTexture = &RDGSourceTexture; + PixelShaderParameters->InputSampler = TStaticSamplerState::GetRHI(); + PixelShaderParameters->RenderTargets[0] = DestinationRenderTarget.GetRenderTargetBinding(); + + ClearUnusedGraphResources(PixelShader, PixelShaderParameters); + + FRHIBlendState* BlendState = FScreenPassPipelineState::FDefaultBlendState::GetRHI(); + FRHIDepthStencilState* DepthStencilState = FScreenPassPipelineState::FDefaultDepthStencilState::GetRHI(); + + const FScreenPassPipelineState PipelineState(VertexShader, PixelShader, BlendState, DepthStencilState); + + GraphBuilder.AddPass( + RDG_EVENT_NAME("RescalingTextureCopy"), + PixelShaderParameters, + ERDGPassFlags::Raster, + [PipelineState, Extent = DestinationRenderTarget.Texture->Desc.Extent, PixelShader, PixelShaderParameters](FRHICommandList& RHICmdList) + { + PipelineState.Validate(); + RHICmdList.SetViewport(0.0f, 0.0f, 0.0f, Extent.X, Extent.Y, 1.0f); + SetScreenPassPipelineState(RHICmdList, PipelineState); + SetShaderParameters(RHICmdList, PixelShader, PixelShader.GetPixelShader(), *PixelShaderParameters); + DrawRectangle( + RHICmdList, + 0, 0, Extent.X, Extent.Y, + 0, 0, Extent.X, Extent.Y, + Extent, + Extent, + PipelineState.VertexShader, + EDRF_UseTriangleOptimization); + }); +} + +void FStyleTransferSceneViewExtension::SubscribeToPostProcessingPass(EPostProcessingPass PassId, FAfterPassCallbackDelegateArray& InOutPassCallbacks, bool bIsPassEnabled) { if (PassId == EPostProcessingPass::Tonemap) { @@ -60,8 +108,88 @@ void FStyleTransferSceneViewExtension::SubscribeToPostProcessingPass(EPostProces } } -FScreenPassTexture FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap_RenderThread( - FRDGBuilder& GraphBuilder, const FSceneView& View, const FPostProcessMaterialInputs& InOutInputs) +FRDGTexture* FStyleTransferSceneViewExtension::TensorToTexture(FRDGBuilder& GraphBuilder, const FRDGTextureDesc& BaseDestinationDesc, const FNeuralTensor& SourceTensor) +{ + FIntVector SourceTensorDimensions = { + CastNarrowingSafe(SourceTensor.GetSize(1)), + CastNarrowingSafe(SourceTensor.GetSize(2)), + CastNarrowingSafe(SourceTensor.GetSize(3)), + }; + + // Reusing the same output description for our back buffer as SceneColor + FRDGTextureDesc DestinationDesc = BaseDestinationDesc; + // this is flipped because the Output tensor has the vertical dimension first + // while unreal has the horizontal dimension first + DestinationDesc.Extent = {SourceTensorDimensions[1], SourceTensorDimensions[0]}; + DestinationDesc.Flags |= TexCreate_RenderTargetable | TexCreate_UAV; + FLinearColor ClearColor(0., 0., 0., 0.); + DestinationDesc.ClearValue = FClearValueBinding(ClearColor); + FRDGTexture* OutputTexture = GraphBuilder.CreateTexture( + DestinationDesc, TEXT("OutputTexture")); + + FRDGBufferRef SourceTensorBuffer = GraphBuilder.RegisterExternalBuffer(SourceTensor.GetPooledBuffer()); + + auto OutputTensorToSceneColorParameters = GraphBuilder.AllocParameters(); + OutputTensorToSceneColorParameters->InputTensor = GraphBuilder.CreateSRV(SourceTensorBuffer, EPixelFormat::PF_R32_FLOAT); + OutputTensorToSceneColorParameters->OutputTexture = GraphBuilder.CreateUAV(OutputTexture); + OutputTensorToSceneColorParameters->TensorVolume = SourceTensor.Num(); + OutputTensorToSceneColorParameters->TextureSize = DestinationDesc.Extent; + FIntVector OutputTensorToSceneColorGroupCount = FComputeShaderUtils::GetGroupCount( + {SourceTensorDimensions.X, SourceTensorDimensions.Y, 1}, + FOutputTensorToSceneColorCS::ThreadGroupSize + ); + + TShaderMapRef OutputTensorToSceneColorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel)); + GraphBuilder.AddPass( + RDG_EVENT_NAME("TensorToTexture"), + OutputTensorToSceneColorParameters, + ERDGPassFlags::Compute, + [OutputTensorToSceneColorCS, OutputTensorToSceneColorParameters, OutputTensorToSceneColorGroupCount](FRHICommandList& RHICommandList) + { + FComputeShaderUtils::Dispatch(RHICommandList, OutputTensorToSceneColorCS, + *OutputTensorToSceneColorParameters, OutputTensorToSceneColorGroupCount); + } + ); + + return OutputTexture; +} + +void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder, const FScreenPassTexture& 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; + + FRDGBufferRef StyleTransferContentInputBuffer = GraphBuilder.RegisterExternalBuffer(DestinationTensor.GetPooledBuffer()); + auto SceneColorToInputTensorParameters = GraphBuilder.AllocParameters(); + SceneColorToInputTensorParameters->TensorVolume = CastNarrowingSafe(DestinationTensor.Num()); + SceneColorToInputTensorParameters->InputTexture = SourceTexture.Texture; + SceneColorToInputTensorParameters->InputTextureSampler = TStaticSamplerState::GetRHI(); + SceneColorToInputTensorParameters->OutputUAV = GraphBuilder.CreateUAV(StyleTransferContentInputBuffer); + SceneColorToInputTensorParameters->OutputDimensions = {InputTensorDimensions.X, InputTensorDimensions.Y}; + SceneColorToInputTensorParameters->HalfPixelUV = FVector2f(0.5f / SceneColorRenderTargetDimensions.X, 0.5 / SceneColorRenderTargetDimensions.Y); + FIntVector SceneColorToInputTensorGroupCount = FComputeShaderUtils::GetGroupCount( + {InputTensorDimensions.X, InputTensorDimensions.Y, 1}, + FSceneColorToInputTensorCS::ThreadGroupSize + ); + + TShaderMapRef SceneColorToInputTensorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel)); + GraphBuilder.AddPass( + RDG_EVENT_NAME("TextureToTensor"), + SceneColorToInputTensorParameters, + ERDGPassFlags::Compute, + [SceneColorToInputTensorCS, SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount](FRHICommandList& RHICommandList) + { + FComputeShaderUtils::Dispatch(RHICommandList, SceneColorToInputTensorCS, + *SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount); + } + ); +} + +FScreenPassTexture FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap_RenderThread(FRDGBuilder& GraphBuilder, const FSceneView& View, const FPostProcessMaterialInputs& InOutInputs) { const FSceneViewFamily& ViewFamily = *View.Family; @@ -79,103 +207,29 @@ FScreenPassTexture FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap RDG_EVENT_SCOPE(GraphBuilder, "StyleTransfer"); - //Get input and output viewports. Backbuffer could be targeting a different region than input viewport - const FScreenPassTextureViewport SceneColorViewport(SceneColor); - - FScreenPassRenderTarget SceneColorRenderTarget(SceneColor, ERenderTargetLoadAction::ELoad); - checkSlow(View.bIsViewInfo); const FViewInfo& ViewInfo = static_cast(View); - /*AddDrawScreenPass(GraphBuilder, RDG_EVENT_NAME("ProcessOCIOColorSpaceXfrm"), ViewInfo, BackBufferViewport, - SceneColorViewport, OCIOPixelShader, Parameters);*/ + const FNeuralTensor& StyleTransferContentInputTensor = StyleTransferNetwork->GetInputTensorForContext(*InferenceContext, 0); - const FNeuralTensor& StyleTransferContentInputTensor = StyleTransferNetwork->GetInputTensorForContext(InferenceContext, 0); + TextureToTensor(GraphBuilder, SceneColor, StyleTransferContentInputTensor); - const FIntVector InputTensorDimensions = { - CastNarrowingSafe(StyleTransferContentInputTensor.GetSize(1)), - CastNarrowingSafe(StyleTransferContentInputTensor.GetSize(2)), - CastNarrowingSafe(StyleTransferContentInputTensor.GetSize(3)), - }; - const FIntPoint SceneColorRenderTargetDimensions = SceneColorRenderTarget.Texture->Desc.Extent; + StyleTransferNetwork->Run(GraphBuilder, *InferenceContext); - FRDGBufferRef StyleTransferContentInputBuffer = GraphBuilder.RegisterExternalBuffer(StyleTransferContentInputTensor.GetPooledBuffer()); - auto SceneColorToInputTensorParameters = GraphBuilder.AllocParameters(); - SceneColorToInputTensorParameters->TensorVolume = CastNarrowingSafe(StyleTransferContentInputTensor.Num()); - SceneColorToInputTensorParameters->InputTexture = SceneColorRenderTarget.Texture; - SceneColorToInputTensorParameters->InputTextureSampler = TStaticSamplerState::GetRHI(); - SceneColorToInputTensorParameters->OutputUAV = GraphBuilder.CreateUAV(StyleTransferContentInputBuffer); - SceneColorToInputTensorParameters->OutputDimensions = {InputTensorDimensions.X, InputTensorDimensions.Y}; - SceneColorToInputTensorParameters->HalfPixelUV = FVector2f(0.5f / SceneColorRenderTargetDimensions.X, 0.5 / SceneColorRenderTargetDimensions.Y); - FIntVector SceneColorToInputTensorGroupCount = FComputeShaderUtils::GetGroupCount( - {InputTensorDimensions.X, InputTensorDimensions.Y, 1}, - FSceneColorToInputTensorCS::ThreadGroupSize - ); + const FNeuralTensor& StyleTransferContentOutputTensor = StyleTransferNetwork->GetInputTensorForContext(*InferenceContext, 0); + FRDGTexture* StyleTransferRenderTargetTexture = TensorToTexture(GraphBuilder, SceneColor.Texture->Desc, StyleTransferContentOutputTensor); - TShaderMapRef SceneColorToInputTensorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel)); - GraphBuilder.AddPass( - RDG_EVENT_NAME("SceneColorToInputTensor"), - SceneColorToInputTensorParameters, - ERDGPassFlags::Compute, - [SceneColorToInputTensorCS, SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount](FRHICommandList& RHICommandList) - { - FComputeShaderUtils::Dispatch(RHICommandList, SceneColorToInputTensorCS, - *SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount); - } - ); - - - const FNeuralTensor& StyleTransferOutputTensor = StyleTransferNetwork->GetOutputTensorForContext(InferenceContext, 0); - FIntVector OutputTensorDimensions = { - CastNarrowingSafe(StyleTransferOutputTensor.GetSize(1)), - CastNarrowingSafe(StyleTransferOutputTensor.GetSize(2)), - CastNarrowingSafe(StyleTransferOutputTensor.GetSize(3)), - }; - // Reusing the same output description for our back buffer as SceneColor - FRDGTextureDesc OutputDesc = SceneColor.Texture->Desc; - // this is flipped because the Output tensor has the vertical dimension first - // while unreal has the horizontal dimension first - OutputDesc.Extent = {OutputTensorDimensions[1], OutputTensorDimensions[0]}; - OutputDesc.Flags |= TexCreate_RenderTargetable | TexCreate_UAV; - FLinearColor ClearColor(0., 0., 0., 0.); - OutputDesc.ClearValue = FClearValueBinding(ClearColor); - FRDGTexture* StyleTransferRenderTargetTexture = GraphBuilder.CreateTexture( - OutputDesc, TEXT("StyleTransferRenderTargetTexture")); TSharedPtr StyleTransferOutputTarget = MakeShared(StyleTransferRenderTargetTexture, SceneColor.ViewRect, ERenderTargetLoadAction::EClear); - StyleTransferNetwork->Run(GraphBuilder, InferenceContext); - - FRDGBufferRef StyleTransferOutputBuffer = GraphBuilder.RegisterExternalBuffer(StyleTransferContentInputTensor.GetPooledBuffer()); - - auto OutputTensorToSceneColorParameters = GraphBuilder.AllocParameters(); - OutputTensorToSceneColorParameters->InputTensor = GraphBuilder.CreateSRV(StyleTransferOutputBuffer, EPixelFormat::PF_FloatRGB); - OutputTensorToSceneColorParameters->OutputTexture = GraphBuilder.CreateUAV(StyleTransferRenderTargetTexture); - FIntVector OutputTensorToSceneColorGroupCount = FComputeShaderUtils::GetGroupCount( - {OutputTensorDimensions.X, OutputTensorDimensions.Y, 1}, - FOutputTensorToSceneColorCS::ThreadGroupSize - ); - - TShaderMapRef OutputTensorToSceneColorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel)); - GraphBuilder.AddPass( - RDG_EVENT_NAME("OutputTensorToSceneColor"), - OutputTensorToSceneColorParameters, - ERDGPassFlags::Compute, - [OutputTensorToSceneColorCS, OutputTensorToSceneColorParameters, OutputTensorToSceneColorGroupCount](FRHICommandList& RHICommandList) - { - FComputeShaderUtils::Dispatch(RHICommandList, OutputTensorToSceneColorCS, - *OutputTensorToSceneColorParameters, OutputTensorToSceneColorGroupCount); - } - ); - TSharedPtr BackBufferRenderTarget; // If the override output is provided it means that this is the last pass in post processing. if (InOutInputs.OverrideOutput.IsValid()) { BackBufferRenderTarget = MakeShared(InOutInputs.OverrideOutput); - // @todo: do not use copy. Resample the styled 1920x960 texture to the fullscreen texture by drawing into the texture - AddCopyTexturePass(GraphBuilder, StyleTransferRenderTargetTexture, BackBufferRenderTarget->Texture); + + AddRescalingTextureCopy(GraphBuilder, *StyleTransferOutputTarget->Texture, *BackBufferRenderTarget); } else { diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp index 6ec57390..cf85574e 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp @@ -5,12 +5,116 @@ #include "NeuralNetwork.h" #include "RenderGraphUtils.h" +#include "StyleTransferModule.h" #include "StyleTransferSceneViewExtension.h" +TAutoConsoleVariable CVarStyleTransferEnabled( + TEXT("r.StyleTransfer.Enabled"), + false, + TEXT("Set to true to enable style transfer") +); + + void UStyleTransferSubsystem::Initialize(FSubsystemCollectionBase& Collection) { Super::Initialize(Collection); + CVarStyleTransferEnabled->OnChangedDelegate().AddUObject(this, &UStyleTransferSubsystem::HandleConsoleVariableChanged); +} + +void UStyleTransferSubsystem::Deinitialize() +{ + StopStylizingViewport(); + + Super::Deinitialize(); +} + +void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportClient) +{ + if (!StylePredictionNetwork->IsLoaded() || !StyleTransferNetwork->IsLoaded()) + { + UE_LOG(LogStyleTransfer, Error, TEXT("Not all networks were loaded, can not stylize viewport.")); + return; + } + + if (!StyleTransferSceneViewExtension) + { + StylePredictionInferenceContext = StylePredictionNetwork->CreateInferenceContext(); + checkf(StylePredictionInferenceContext != INDEX_NONE, TEXT("Could not create inference context for StylePredictionNetwork")); + StyleTransferInferenceContext = MakeShared(StyleTransferNetwork->CreateInferenceContext()); + checkf(*StyleTransferInferenceContext != INDEX_NONE, TEXT("Could not create inference context for StyleTransferNetwork")); + + FlushRenderingCommands(); + StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension(ViewportClient, StyleTransferNetwork, StyleTransferInferenceContext.ToSharedRef()); + } + StyleTransferSceneViewExtension->SetEnabled(true); +} + +void UStyleTransferSubsystem::StopStylizingViewport() +{ + FlushRenderingCommands(); + StyleTransferSceneViewExtension.Reset(); + if (StylePredictionInferenceContext != INDEX_NONE) + { + StylePredictionNetwork->DestroyInferenceContext(StylePredictionInferenceContext); + } + if (StyleTransferInferenceContext && *StyleTransferInferenceContext != INDEX_NONE) + { + StyleTransferNetwork->DestroyInferenceContext(*StyleTransferInferenceContext); + *StyleTransferInferenceContext = INDEX_NONE; + StyleTransferInferenceContext.Reset(); + } +} + +void UStyleTransferSubsystem::UpdateStyle(const FNeuralTensor& StyleImage) +{ + checkf(StyleTransferSceneViewExtension.IsValid(), TEXT("Can not update style while not stylizing")); + checkf(StyleTransferInferenceContext.IsValid(), TEXT("Can not update style without inference context")); + + StylePredictionNetwork->SetInputFromArrayCopy(StyleImage.GetArrayCopy()); + + ENQUEUE_RENDER_COMMAND(StylePrediction)([this](FRHICommandListImmediate& RHICommandList) + { + FRDGBuilder GraphBuilder(RHICommandList); + + StylePredictionNetwork->Run(GraphBuilder, StylePredictionInferenceContext); + + const FNeuralTensor& OutputStyleParams = StylePredictionNetwork->GetOutputTensorForContext(StylePredictionInferenceContext, 0); + const FNeuralTensor& InputStyleParams = StyleTransferNetwork->GetInputTensorForContext(*StyleTransferInferenceContext, StyleTransferStyleParamsInputIndex); + + 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(); + }); + + FlushRenderingCommands(); +} + +void UStyleTransferSubsystem::HandleConsoleVariableChanged(IConsoleVariable* ConsoleVariable) +{ + check(ConsoleVariable == CVarStyleTransferEnabled.AsVariable()); + + StopStylizingViewport(); + + if (CVarStyleTransferEnabled->GetBool()) + { + if(!(StyleTransferNetwork || StylePredictionNetwork)) + { + LoadNetworks(); + } + StartStylizingViewport(GetGameInstance()->GetGameViewportClient()); + } +} + +void UStyleTransferSubsystem::LoadNetworks() +{ StyleTransferNetwork = LoadObject(this, TEXT("/StyleTransfer/NN_StyleTransfer.NN_StyleTransfer")); StylePredictionNetwork = LoadObject(this, TEXT("/StyleTransfer/NN_StylePredictor.NN_StylePredictor")); @@ -29,7 +133,7 @@ void UStyleTransferSubsystem::Initialize(FSubsystemCollectionBase& Collection) } else { - UE_LOG(LogTemp, Warning, TEXT("StyleTransferNetwork could not be loaded")); + UE_LOG(LogStyleTransfer, Error, TEXT("StyleTransferNetwork could not be loaded")); } @@ -39,61 +143,6 @@ void UStyleTransferSubsystem::Initialize(FSubsystemCollectionBase& Collection) } else { - UE_LOG(LogTemp, Warning, TEXT("StylePredictionNetwork could not be loaded.")); + UE_LOG(LogStyleTransfer, Error, TEXT("StylePredictionNetwork could not be loaded.")); } } - -void UStyleTransferSubsystem::Deinitialize() -{ - StyleTransferSceneViewExtension.Reset(); - if(StylePredictionInferenceContext != INDEX_NONE) - { - StylePredictionNetwork->DestroyInferenceContext(StylePredictionInferenceContext); - } - if(StyleTransferInferenceContext != INDEX_NONE) - { - StyleTransferNetwork->DestroyInferenceContext(StyleTransferInferenceContext); - } - - Super::Deinitialize(); -} - -void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportClient) -{ - StylePredictionInferenceContext = StylePredictionNetwork->CreateInferenceContext(); - checkf(StylePredictionInferenceContext != INDEX_NONE, TEXT("Could not create inference context for StylePredictionNetwork")); - StyleTransferInferenceContext = StyleTransferNetwork->CreateInferenceContext(); - checkf(StyleTransferInferenceContext != INDEX_NONE, TEXT("Could not create inference context for StyleTransferNetwork")); - - StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension(ViewportClient, StyleTransferNetwork, StyleTransferInferenceContext); -} - -void UStyleTransferSubsystem::UpdateStyle(const FNeuralTensor& StyleImage) -{ - checkf(StyleTransferSceneViewExtension.IsValid(), TEXT("Can not update style while not stylizing")); - - StylePredictionNetwork->SetInputFromArrayCopy(StyleImage.GetArrayCopy()); - - ENQUEUE_RENDER_COMMAND(StylePrediction)([this](FRHICommandListImmediate& RHICommandList) - { - FRDGBuilder GraphBuilder(RHICommandList); - - StylePredictionNetwork->Run(GraphBuilder, StylePredictionInferenceContext); - - const FNeuralTensor& OutputStyleParams = StylePredictionNetwork->GetOutputTensorForContext(StylePredictionInferenceContext, 0); - const FNeuralTensor& InputStyleParams = StyleTransferNetwork->GetInputTensorForContext(StyleTransferInferenceContext, StyleTransferStyleParamsInputIndex); - - 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(); - }); - - FlushRenderingCommands(); -} diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h index 6da3daca..1eca292c 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h @@ -1,6 +1,8 @@ #pragma once #include "SceneViewExtension.h" +struct FNeuralTensor; +struct FScreenPassRenderTarget; class UNeuralNetwork; class FStyleTransferSceneViewExtension : public FSceneViewExtensionBase @@ -9,21 +11,40 @@ public: using Ptr = TSharedPtr; using Ref = TSharedRef; - FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, int32 InInferenceContext); + FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef InInferenceContext); // - ISceneViewExtension virtual void SubscribeToPostProcessingPass(EPostProcessingPass Pass, FAfterPassCallbackDelegateArray& InOutPassCallbacks, bool bIsPassEnabled) override; FScreenPassTexture PostProcessPassAfterTonemap_RenderThread(FRDGBuilder& GraphBuilder, const FSceneView& View, const FPostProcessMaterialInputs& InOutInputs); virtual void SetupViewFamily(FSceneViewFamily& InViewFamily) override; - virtual void SetupView(FSceneViewFamily& InViewFamily, FSceneView& InView) override {}; - virtual void BeginRenderViewFamily(FSceneViewFamily& InViewFamily) override {}; + + virtual void SetupView(FSceneViewFamily& InViewFamily, FSceneView& InView) override + { + } + + virtual void BeginRenderViewFamily(FSceneViewFamily& InViewFamily) override + { + } + + virtual bool IsActiveThisFrame_Internal(const FSceneViewExtensionContext& Context) const override; // -- + void SetEnabled(bool bInIsEnabled) { bIsEnabled = bInIsEnabled; } + bool IsEnabled() const { return bIsEnabled; } + private: + /** The actual Network pointer is not tracked so we need a WeakPtr too so we can check its validity on the game thread. */ + TWeakObjectPtr StyleTransferNetworkWeakPtr; TObjectPtr StyleTransferNetwork; FViewportClient* LinkedViewportClient; - int32 InferenceContext = -1; + TSharedRef InferenceContext = MakeShared(-1); + + bool bIsEnabled = true; + + void AddRescalingTextureCopy(FRDGBuilder& GraphBuilder, FRDGTexture& RDGSourceTexture, FScreenPassRenderTarget& DestinationRenderTarget); + FRDGTexture* TensorToTexture(FRDGBuilder& GraphBuilder, const FRDGTextureDesc& BaseDestinationDesc, const FNeuralTensor& SourceTensor); + void TextureToTensor(FRDGBuilder& GraphBuilder, const FScreenPassTexture& SourceTexture, const FNeuralTensor& DestinationTensor); }; diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h index d65be25f..023cdd83 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h @@ -24,6 +24,7 @@ public: // -- void StartStylizingViewport(FViewportClient* ViewportClient); + void StopStylizingViewport(); void UpdateStyle(const FNeuralTensor& StyleImage); @@ -37,7 +38,12 @@ private: TObjectPtr StylePredictionNetwork; int32 StylePredictionInferenceContext = INDEX_NONE; - int32 StyleTransferInferenceContext = INDEX_NONE; + TSharedPtr StyleTransferInferenceContext; + int32 StyleTransferStyleParamsInputIndex = INDEX_NONE; + + void HandleConsoleVariableChanged(IConsoleVariable*); + + void LoadNetworks(); }; diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/OutputTensorToSceneColorCS.h b/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/OutputTensorToSceneColorCS.h index 7c0ff926..8766c0e6 100644 --- a/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/OutputTensorToSceneColorCS.h +++ b/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/OutputTensorToSceneColorCS.h @@ -20,6 +20,7 @@ class STYLETRANSFERSHADERS_API FOutputTensorToSceneColorCS : public FGlobalShade BEGIN_SHADER_PARAMETER_STRUCT(FParameters, ) SHADER_PARAMETER(uint32, TensorVolume) + SHADER_PARAMETER(FIntPoint, TextureSize) SHADER_PARAMETER_RDG_BUFFER_SRV(Buffer, InputTensor) SHADER_PARAMETER_RDG_TEXTURE_UAV(RWTexture2D, OutputTexture) END_SHADER_PARAMETER_STRUCT()