From 97c5d9ed573abb8574965c96cab3a4ef7442d541 Mon Sep 17 00:00:00 2001 From: Manuel Wagner Date: Tue, 23 Aug 2022 17:28:59 +0200 Subject: [PATCH] added style transfer plugin --- Lyra.uproject | 4 + Plugins/StyleTransfer/Resources/Icon128.png | 3 + .../Private/OutputTensorToSceneColor.usf | 30 ++++ .../Private/SceneColorToInputTensor.usf | 27 +++ .../Private/StyleTransferModule.cpp | 23 +++ .../StyleTransferSceneViewExtension.cpp | 158 ++++++++++++++++++ .../Private/StyleTransferSubsystem.cpp | 60 +++++++ .../Public/StyleTransferModule.h | 17 ++ .../Public/StyleTransferSceneViewExtension.h | 29 ++++ .../Public/StyleTransferSubsystem.h | 37 ++++ .../StyleTransfer/StyleTransfer.Build.cs | 56 +++++++ .../Private/OutputTensorToSceneColorCS.cpp | 15 ++ .../Private/SceneColorToInputTensorCS.cpp | 19 +++ .../Private/StyleTransferShaders.cpp | 33 ++++ .../Public/OutputTensorToSceneColorCS.h | 34 ++++ .../Public/SceneColorToInputTensorCS.h | 40 +++++ .../Public/StyleTransferShaders.h | 13 ++ .../StyleTransferShaders.Build.cs | 29 ++++ Plugins/StyleTransfer/StyleTransfer.uplugin | 37 ++++ Source/LyraGame/LyraGame.Build.cs | 3 +- 20 files changed, 666 insertions(+), 1 deletion(-) create mode 100644 Plugins/StyleTransfer/Resources/Icon128.png create mode 100644 Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf create mode 100644 Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf create mode 100644 Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferModule.cpp create mode 100644 Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp create mode 100644 Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp create mode 100644 Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferModule.h create mode 100644 Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h create mode 100644 Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h create mode 100644 Plugins/StyleTransfer/Source/StyleTransfer/StyleTransfer.Build.cs create mode 100644 Plugins/StyleTransfer/Source/StyleTransferShaders/Private/OutputTensorToSceneColorCS.cpp create mode 100644 Plugins/StyleTransfer/Source/StyleTransferShaders/Private/SceneColorToInputTensorCS.cpp create mode 100644 Plugins/StyleTransfer/Source/StyleTransferShaders/Private/StyleTransferShaders.cpp create mode 100644 Plugins/StyleTransfer/Source/StyleTransferShaders/Public/OutputTensorToSceneColorCS.h create mode 100644 Plugins/StyleTransfer/Source/StyleTransferShaders/Public/SceneColorToInputTensorCS.h create mode 100644 Plugins/StyleTransfer/Source/StyleTransferShaders/Public/StyleTransferShaders.h create mode 100644 Plugins/StyleTransfer/Source/StyleTransferShaders/StyleTransferShaders.Build.cs create mode 100644 Plugins/StyleTransfer/StyleTransfer.uplugin diff --git a/Lyra.uproject b/Lyra.uproject index e1697bd0..f5f3e85e 100644 --- a/Lyra.uproject +++ b/Lyra.uproject @@ -285,6 +285,10 @@ { "Name": "TopDownArena", "Enabled": true + }, + { + "Name": "StyleTransfer", + "Enabled": true } ], "EpicSampleNameHash": "451731683" diff --git a/Plugins/StyleTransfer/Resources/Icon128.png b/Plugins/StyleTransfer/Resources/Icon128.png new file mode 100644 index 00000000..26245f6a --- /dev/null +++ b/Plugins/StyleTransfer/Resources/Icon128.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f7239efaeefbd82de33ebe18518e50de075ea4188a468a9e4991396433d2275f +size 12699 diff --git a/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf b/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf new file mode 100644 index 00000000..a301fb59 --- /dev/null +++ b/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf @@ -0,0 +1,30 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#include "/Engine/Public/Platform.ush" + + +RWTexture2D OutputTexture; +Buffer InputTensor; + +[numthreads(1, 1, 1)] +void SceneColorToInputTensorCS(in const uint3 DispatchThreadID : SV_DispatchThreadID) +{ + const uint GlobalIndex = DispatchThreadID.x; + + uint TensorVolume; + InputTensor.GetDimensions(TensorVolume); + + if (GlobalIndex >= TensorVolume) + { + return; + } + uint2 TextureSize = 0; + OutputTexture.GetDimensions(TextureSize.x, TextureSize.y); + const uint PixelIndex = GlobalIndex / 3; + const uint ChannelIndex = GlobalIndex % 3; + const uint MajorIndex = PixelIndex / TextureSize.x; + const uint MinorIndex = PixelIndex % TextureSize.x; + const uint2 TextureCoords = uint2(MajorIndex, MinorIndex); + + OutputTexture[TextureCoords][ChannelIndex] = InputTensor[GlobalIndex]; +} diff --git a/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf b/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf new file mode 100644 index 00000000..72ce6ffe --- /dev/null +++ b/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf @@ -0,0 +1,27 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#include "/Engine/Public/Platform.ush" + +Texture2D InputTexture; +RWBuffer OutputUAV; + +[numthreads(THREADGROUP_SIZE_X, THREADGROUP_SIZE_Y, THREADGROUP_SIZE_Z)] +void SceneColorToInputTensorCS(in const uint3 DispatchThreadID : SV_DispatchThreadID) +{ + uint2 TexelCoordinate = DispatchThreadID.xy; + + const uint GlobalIndex = TexelCoordinate.x * TexelCoordinate.y * 3; + + uint TensorVolume; + OutputUAV.GetDimensions(TensorVolume); + + if (GlobalIndex >= TensorVolume) + { + return; + } + + float4 TextureValue = InputTexture[DispatchThreadID.xy]; + OutputUAV[GlobalIndex + 0] = TextureValue.r; + OutputUAV[GlobalIndex + 1] = TextureValue.g; + OutputUAV[GlobalIndex + 2] = TextureValue.b; +} diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferModule.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferModule.cpp new file mode 100644 index 00000000..8623dfe7 --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferModule.cpp @@ -0,0 +1,23 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#include "StyleTransferModule.h" + +#include "ShaderCore.h" +#include "Interfaces/IPluginManager.h" +#include "Logging/LogMacros.h" + +DEFINE_LOG_CATEGORY(LogStyleTransfer) + +#define LOCTEXT_NAMESPACE "FStyleTransferModule" + +void FStyleTransferModule::StartupModule() +{ +} + +void FStyleTransferModule::ShutdownModule() +{ +} + +#undef LOCTEXT_NAMESPACE + +IMPLEMENT_MODULE(FStyleTransferModule, StyleTransfer) diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp new file mode 100644 index 00000000..96d474b3 --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp @@ -0,0 +1,158 @@ +#include "StyleTransferSceneViewExtension.h" + +#include "CoreGlobals.h" +#include "SceneView.h" + +#include "GlobalShader.h" +#include "PipelineStateCache.h" +#include "RenderTargetPool.h" +#include "Shader.h" + +#include "RHI.h" +#include "SceneView.h" +#include "ScreenPass.h" +#include "CommonRenderResources.h" +#include "NeuralNetwork.h" +#include "RenderGraphEvent.h" +#include "PostProcess/PostProcessing.h" +#include "Containers/DynamicRHIResourceArray.h" +#include "PostProcess/PostProcessMaterial.h" +#include "OutputTensorToSceneColorCS.h" +#include "SceneColorToInputTensorCS.h" + +template +OutType CastNarrowingSafe(InType InValue) +{ + if (InValue > TNumericLimits::Max()) + { + return TNumericLimits::Max(); + } + if (InValue < TNumericLimits::Min()) + { + return TNumericLimits::Min(); + } + return static_cast(InValue); +} + + +FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork) + : FSceneViewExtensionBase(AutoRegister) + , StyleTransferNetwork(InStyleTransferNetwork) + , LinkedViewportClient(AssociatedViewportClient) +{ + ensure(InStyleTransferNetwork->GetDeviceType() == ENeuralDeviceType::GPU); +} + +void FStyleTransferSceneViewExtension::SetupViewFamily(FSceneViewFamily& InViewFamily) +{ + InferenceContext = StyleTransferNetwork->CreateInferenceContext(); +} + +void FStyleTransferSceneViewExtension::SubscribeToPostProcessingPass(EPostProcessingPass PassId, + FAfterPassCallbackDelegateArray& + InOutPassCallbacks, bool bIsPassEnabled) +{ + if (PassId == EPostProcessingPass::Tonemap) + { + InOutPassCallbacks.Add( + FAfterPassCallbackDelegate::CreateRaw( + this, &FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap_RenderThread)); + } +} + +FScreenPassTexture FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap_RenderThread( + FRDGBuilder& GraphBuilder, const FSceneView& View, const FPostProcessMaterialInputs& InOutInputs) +{ + const FSceneViewFamily& ViewFamily = *View.Family; + + const FScreenPassTexture& SceneColor = InOutInputs.Textures[(uint32)EPostProcessMaterialInput::SceneColor]; + + if (!EnumHasAnyFlags(SceneColor.Texture->Desc.Flags, TexCreate_ShaderResource)) + { + return SceneColor; + } + + if (!SceneColor.IsValid()) + { + return SceneColor; + } + + RDG_EVENT_SCOPE(GraphBuilder, "StyleTransfer"); + + FScreenPassRenderTarget BackBufferRenderTarget; + + // If the override output is provided it means that this is the last pass in post processing. + if (InOutInputs.OverrideOutput.IsValid()) + { + BackBufferRenderTarget = InOutInputs.OverrideOutput; + } + else + { + // Reusing the same output description for our back buffer as SceneColor when it's not overriden + FRDGTextureDesc OutputDesc = SceneColor.Texture->Desc; + OutputDesc.Flags |= TexCreate_RenderTargetable; + FLinearColor ClearColor(0., 0., 0., 0.); + OutputDesc.ClearValue = FClearValueBinding(ClearColor); + + FRDGTexture* BackBufferRenderTargetTexture = GraphBuilder.CreateTexture( + OutputDesc, TEXT("BackBufferRenderTargetTexture")); + BackBufferRenderTarget = FScreenPassRenderTarget(BackBufferRenderTargetTexture, SceneColor.ViewRect, + ERenderTargetLoadAction::EClear); + } + + //Get input and output viewports. Backbuffer could be targeting a different region than input viewport + const FScreenPassTextureViewport SceneColorViewport(SceneColor); + const FScreenPassTextureViewport BackBufferViewport(BackBufferRenderTarget); + + 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);*/ + + auto SceneColorToInputTensorParameters = GraphBuilder.AllocParameters(); + FNeuralTensor InputTensor = StyleTransferNetwork->GetInputTensor(); + SceneColorToInputTensorParameters->TensorVolume = CastNarrowingSafe(InputTensor.Num()); + SceneColorToInputTensorParameters->InputTexture = SceneColorRenderTarget.Texture; + SceneColorToInputTensorParameters->OutputUAV = InputTensor.GetBufferUAVRef(); + FIntVector SceneColorToInputTensorGroupCount; + + TShaderMapRef SceneColorToInputTensorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel)); + ClearUnusedGraphResources(SceneColorToInputTensorCS, SceneColorToInputTensorParameters); + GraphBuilder.AddPass( + RDG_EVENT_NAME("SceneColorToInputTensor"), + SceneColorToInputTensorParameters, + ERDGPassFlags::Compute, + [SceneColorToInputTensorCS, SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount](FRHICommandList& RHICommandList) + { + FComputeShaderUtils::Dispatch(RHICommandList, SceneColorToInputTensorCS, + *SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount); + } + ); + + StyleTransferNetwork->GetInputDataPointerMutableForContext(InferenceContext, 0); + StyleTransferNetwork->Run(GraphBuilder, InferenceContext); + + auto OutputTensorToSceneColorParameters = GraphBuilder.AllocParameters(); + FNeuralTensor OutputTensor = StyleTransferNetwork->GetOutputTensor(0); + OutputTensorToSceneColorParameters->InputTensor = OutputTensor.GetBufferSRVRef(); + OutputTensorToSceneColorParameters->OutputTexture = GraphBuilder.CreateUAV(SceneColorRenderTarget.Texture); + FIntVector OutputTensorToSceneColorGroupCount; + + TShaderMapRef OutputTensorToSceneColorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel)); + ClearUnusedGraphResources(OutputTensorToSceneColorCS, OutputTensorToSceneColorParameters); + GraphBuilder.AddPass( + RDG_EVENT_NAME("OutputTensorToSceneColor"), + OutputTensorToSceneColorParameters, + ERDGPassFlags::Compute, + [OutputTensorToSceneColorCS, OutputTensorToSceneColorParameters, OutputTensorToSceneColorGroupCount](FRHICommandList& RHICommandList) + { + FComputeShaderUtils::Dispatch(RHICommandList, OutputTensorToSceneColorCS, + *OutputTensorToSceneColorParameters, OutputTensorToSceneColorGroupCount); + } + ); + + return MoveTemp(BackBufferRenderTarget); +} diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp new file mode 100644 index 00000000..aff9bd39 --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp @@ -0,0 +1,60 @@ +// Fill out your copyright notice in the Description page of Project Settings. + + +#include "StyleTransferSubsystem.h" + +#include "NeuralNetwork.h" +#include "StyleTransferSceneViewExtension.h" + +void UStyleTransferSubsystem::Initialize(FSubsystemCollectionBase& Collection) +{ + Super::Initialize(Collection); + + StyleTransferNetwork = NewObject(); + + FString ONNXModelFilePath = TEXT("SOME_PARENT_FOLDER/SOME_ONNX_FILE_NAME.onnx"); + if (StyleTransferNetwork->Load(ONNXModelFilePath)) + { + if (StyleTransferNetwork->IsGPUSupported()) + { + StyleTransferNetwork->SetDeviceType(ENeuralDeviceType::GPU); + } + else + { + StyleTransferNetwork->SetDeviceType(ENeuralDeviceType::CPU); + } + } + else + { + UE_LOG(LogTemp, Warning, TEXT("StyleTransferNetwork could not loaded from %s."), *ONNXModelFilePath); + } + + StylePredictionNetwork = NewObject(); + + ONNXModelFilePath = TEXT("SOME_PARENT_FOLDER/SOME_ONNX_FILE_NAME.onnx"); + if (StylePredictionNetwork->Load(ONNXModelFilePath)) + { + if (StylePredictionNetwork->IsGPUSupported()) + { + StylePredictionNetwork->SetDeviceType(ENeuralDeviceType::GPU); + } + else + { + StylePredictionNetwork->SetDeviceType(ENeuralDeviceType::CPU); + } + } + else + { + UE_LOG(LogTemp, Warning, TEXT("StylePredictionNetwork could not loaded from %s."), *ONNXModelFilePath); + } +} + +void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportClient) +{ + StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension(ViewportClient, StyleTransferNetwork); +} + +void UStyleTransferSubsystem::UpdateStyle(FNeuralTensor StyleImage) +{ + +} diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferModule.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferModule.h new file mode 100644 index 00000000..20a73f5a --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferModule.h @@ -0,0 +1,17 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#pragma once + +#include "CoreMinimal.h" +#include "Modules/ModuleManager.h" + +DECLARE_LOG_CATEGORY_EXTERN(LogStyleTransfer, Log, All); + +class FStyleTransferModule : public IModuleInterface +{ +public: + + /** IModuleInterface implementation */ + virtual void StartupModule() override; + virtual void ShutdownModule() override; +}; diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h new file mode 100644 index 00000000..5c9f3038 --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h @@ -0,0 +1,29 @@ +#pragma once +#include "SceneViewExtension.h" + +class UNeuralNetwork; + +class FStyleTransferSceneViewExtension : public FSceneViewExtensionBase +{ +public: + using Ptr = TSharedPtr; + using Ref = TSharedRef; + + FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork); + + // - 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 {}; + // -- + +private: + TObjectPtr StyleTransferNetwork; + + FViewportClient* LinkedViewportClient; + + int32 InferenceContext = -1; +}; diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h new file mode 100644 index 00000000..6436c3b7 --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h @@ -0,0 +1,37 @@ +// Fill out your copyright notice in the Description page of Project Settings. + +#pragma once + +#include "CoreMinimal.h" +#include "NeuralTensor.h" +#include "StyleTransferSceneViewExtension.h" +#include "Subsystems/GameInstanceSubsystem.h" +#include "UObject/Object.h" +#include "StyleTransferSubsystem.generated.h" + +/** + * + */ +UCLASS() +class STYLETRANSFER_API UStyleTransferSubsystem : public UGameInstanceSubsystem +{ + GENERATED_BODY() + +public: + // - UGameInstanceSubsystem + virtual void Initialize(FSubsystemCollectionBase& Collection) override; + // -- + + void StartStylizingViewport(FViewportClient* ViewportClient); + + void UpdateStyle(FNeuralTensor StyleImage); + +private: + FStyleTransferSceneViewExtension::Ptr StyleTransferSceneViewExtension; + + UPROPERTY() + TObjectPtr StyleTransferNetwork; + + UPROPERTY() + TObjectPtr StylePredictionNetwork; +}; diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/StyleTransfer.Build.cs b/Plugins/StyleTransfer/Source/StyleTransfer/StyleTransfer.Build.cs new file mode 100644 index 00000000..e67641e4 --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransfer/StyleTransfer.Build.cs @@ -0,0 +1,56 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +using System.IO; +using UnrealBuildTool; + +public class StyleTransfer : ModuleRules +{ + public StyleTransfer(ReadOnlyTargetRules Target) : base(Target) + { + PCHUsage = ModuleRules.PCHUsageMode.UseExplicitOrSharedPCHs; + + PublicIncludePaths.AddRange( + new string[] + { + } + ); + + string EngineDir = Path.GetFullPath(Target.RelativeEnginePath); + PrivateIncludePaths.AddRange( + new string[] + { + Path.Combine(EngineDir, "Source/Runtime/Renderer/Private"), + } + ); + + + PublicDependencyModuleNames.AddRange( + new string[] + { + "Core", + "NeuralNetworkInference", + } + ); + + + PrivateDependencyModuleNames.AddRange( + new string[] + { + "CoreUObject", + "Engine", + "RHI", + "RenderCore", + "Renderer", + "Projects", + "StyleTransferShaders" + } + ); + + + DynamicallyLoadedModuleNames.AddRange( + new string[] + { + } + ); + } +} \ No newline at end of file diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/OutputTensorToSceneColorCS.cpp b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/OutputTensorToSceneColorCS.cpp new file mode 100644 index 00000000..f9791020 --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/OutputTensorToSceneColorCS.cpp @@ -0,0 +1,15 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#include "OutputTensorToSceneColorCS.h" +#include "Utils.h" + + + +void FOutputTensorToSceneColorCS::ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment) +{ + FGlobalShader::ModifyCompilationEnvironment(Parameters, OutEnvironment); +} + +IMPLEMENT_GLOBAL_SHADER(FOutputTensorToSceneColorCS, + "/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf", + "OutputTensorToSceneColorCS", SF_Compute); // Path defined in StyleTransferModule.cpp diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/SceneColorToInputTensorCS.cpp b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/SceneColorToInputTensorCS.cpp new file mode 100644 index 00000000..0252df9d --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/SceneColorToInputTensorCS.cpp @@ -0,0 +1,19 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#include "SceneColorToInputTensorCS.h" +#include "Utils.h" + + + +void FSceneColorToInputTensorCS::ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment) +{ + FGlobalShader::ModifyCompilationEnvironment(Parameters, OutEnvironment); + + OutEnvironment.SetDefine(TEXT("THREADGROUP_SIZE_X"), THREADGROUP_SIZE_X); + OutEnvironment.SetDefine(TEXT("THREADGROUP_SIZE_Y"), THREADGROUP_SIZE_Y); + OutEnvironment.SetDefine(TEXT("THREADGROUP_SIZE_Z"), 1); +} + +IMPLEMENT_GLOBAL_SHADER(FSceneColorToInputTensorCS, + "/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf", + "SceneColorToInputTensorCS", SF_Compute); // Path defined in StyleTransferModule.cpp diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/StyleTransferShaders.cpp b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/StyleTransferShaders.cpp new file mode 100644 index 00000000..4192487b --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/StyleTransferShaders.cpp @@ -0,0 +1,33 @@ +#include "StyleTransferShaders.h" + +#include "Interfaces/IPluginManager.h" + +DEFINE_LOG_CATEGORY(LogStyleTransferShaders); + +#define LOCTEXT_NAMESPACE "FStyleTransferShadersModule" + +void FStyleTransferShadersModule::StartupModule() +{ + // Add shader directory + const TSharedPtr StyleTransferPlugin = IPluginManager::Get().FindPlugin(TEXT("StyleTransfer")); + if (StyleTransferPlugin.IsValid()) + { + const FString RealShaderDirectory = StyleTransferPlugin->GetBaseDir() / TEXT("Shaders"); // TEXT("../../../Plugins/StyleTransfer/Shaders"); + const FString VirtualShaderDirectory = TEXT("/Plugins/StyleTransfer/Shaders"); + AddShaderSourceDirectoryMapping(VirtualShaderDirectory, RealShaderDirectory); + } + else + { + UE_LOG(LogStyleTransferShaders, Warning, + TEXT("FStyleTransferModule::StartupModule(): StyleTransferPlugin was nullptr, shaders directory not added.")); + } +} + +void FStyleTransferShadersModule::ShutdownModule() +{ + +} + +#undef LOCTEXT_NAMESPACE + +IMPLEMENT_MODULE(FStyleTransferShadersModule, StyleTransferShaders) \ No newline at end of file diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/OutputTensorToSceneColorCS.h b/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/OutputTensorToSceneColorCS.h new file mode 100644 index 00000000..787b4765 --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/OutputTensorToSceneColorCS.h @@ -0,0 +1,34 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#pragma once + +// GPU/RHI/shaders +#include "GlobalShader.h" +#include "RHI.h" +#include "ProfilingDebugging/RealtimeGPUProfiler.h" +#include "RenderGraphUtils.h" +#include "ShaderParameterUtils.h" + + + +class STYLETRANSFERSHADERS_API FOutputTensorToSceneColorCS : public FGlobalShader +{ + DECLARE_GLOBAL_SHADER(FOutputTensorToSceneColorCS); + SHADER_USE_PARAMETER_STRUCT(FOutputTensorToSceneColorCS, FGlobalShader) + + BEGIN_SHADER_PARAMETER_STRUCT(FParameters, ) + SHADER_PARAMETER(uint32, TensorVolume) + SHADER_PARAMETER_RDG_BUFFER_SRV(Buffer, InputTensor) + SHADER_PARAMETER_RDG_TEXTURE_UAV(RWTexture2D, OutputTexture) + END_SHADER_PARAMETER_STRUCT() + + // - FShader + static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters) { + return GetMaxSupportedFeatureLevel(Parameters.Platform) >= ERHIFeatureLevel::SM5; + } + + static void ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment); + // -- + +private: +}; diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/SceneColorToInputTensorCS.h b/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/SceneColorToInputTensorCS.h new file mode 100644 index 00000000..0697145c --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/SceneColorToInputTensorCS.h @@ -0,0 +1,40 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#pragma once + +// GPU/RHI/shaders +#include "GlobalShader.h" +#include "RHI.h" +#include "ProfilingDebugging/RealtimeGPUProfiler.h" +#include "RenderGraphUtils.h" +#include "ShaderParameterUtils.h" + + + +class STYLETRANSFERSHADERS_API FSceneColorToInputTensorCS : public FGlobalShader +{ + DECLARE_GLOBAL_SHADER(FSceneColorToInputTensorCS); + SHADER_USE_PARAMETER_STRUCT(FSceneColorToInputTensorCS, FGlobalShader) + + static const uint32 THREADGROUP_SIZE_X = 8; + static const uint32 THREADGROUP_SIZE_Y = 8; + + BEGIN_SHADER_PARAMETER_STRUCT(FParameters, ) + // Input variables + SHADER_PARAMETER(uint32, TensorVolume) + // SRV/UAV variables + SHADER_PARAMETER_RDG_BUFFER_UAV(RWBuffer, OutputUAV) + // Optional SRV/UAV variables + SHADER_PARAMETER_RDG_TEXTURE(Texture2D, InputTexture) + END_SHADER_PARAMETER_STRUCT() + + // - FShader + static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters) { + return GetMaxSupportedFeatureLevel(Parameters.Platform) >= ERHIFeatureLevel::SM5; + } + + static void ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment); + // -- + +private: +}; diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/StyleTransferShaders.h b/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/StyleTransferShaders.h new file mode 100644 index 00000000..336015fc --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/StyleTransferShaders.h @@ -0,0 +1,13 @@ +#pragma once + +#include "CoreMinimal.h" +#include "Modules/ModuleManager.h" + +DECLARE_LOG_CATEGORY_EXTERN(LogStyleTransferShaders, Log, All); + +class FStyleTransferShadersModule : public IModuleInterface +{ +public: + virtual void StartupModule() override; + virtual void ShutdownModule() override; +}; diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/StyleTransferShaders.Build.cs b/Plugins/StyleTransfer/Source/StyleTransferShaders/StyleTransferShaders.Build.cs new file mode 100644 index 00000000..3e385675 --- /dev/null +++ b/Plugins/StyleTransfer/Source/StyleTransferShaders/StyleTransferShaders.Build.cs @@ -0,0 +1,29 @@ +using UnrealBuildTool; + +public class StyleTransferShaders : ModuleRules +{ + public StyleTransferShaders(ReadOnlyTargetRules Target) : base(Target) + { + PCHUsage = ModuleRules.PCHUsageMode.UseExplicitOrSharedPCHs; + + PublicDependencyModuleNames.AddRange( + new string[] + { + "Core", + "RenderCore", + } + ); + + PrivateDependencyModuleNames.AddRange( + new string[] + { + "CoreUObject", + "Engine", + "Slate", + "SlateCore", + "RHI", + "Projects" + } + ); + } +} \ No newline at end of file diff --git a/Plugins/StyleTransfer/StyleTransfer.uplugin b/Plugins/StyleTransfer/StyleTransfer.uplugin new file mode 100644 index 00000000..e9798c49 --- /dev/null +++ b/Plugins/StyleTransfer/StyleTransfer.uplugin @@ -0,0 +1,37 @@ +{ + "FileVersion": 3, + "Version": 1, + "VersionName": "1.0", + "FriendlyName": "StyleTransfer", + "Description": "Neural Style Transfer for rendering in different styles", + "Category": "Other", + "CreatedBy": "Manuel Wagner", + "CreatedByURL": "https://singinwhale.com", + "DocsURL": "", + "MarketplaceURL": "", + "CanContainContent": true, + "IsBetaVersion": false, + "IsExperimentalVersion": false, + "Installed": false, + "Modules": [ + { + "Name": "StyleTransfer", + "Type": "Runtime", + "LoadingPhase": "PostConfigInit", + "WhitelistPlatforms": [ + "Win64" + ] + }, + { + "Name": "StyleTransferShaders", + "Type": "Runtime", + "LoadingPhase": "PostConfigInit" + } + ], + "Plugins": [ + { + "Name": "NeuralNetworkInference", + "Enabled": true + } + ] +} \ No newline at end of file diff --git a/Source/LyraGame/LyraGame.Build.cs b/Source/LyraGame/LyraGame.Build.cs index 13d720ce..00faed07 100644 --- a/Source/LyraGame/LyraGame.Build.cs +++ b/Source/LyraGame/LyraGame.Build.cs @@ -41,7 +41,8 @@ public class LyraGame : ModuleRules "CommonLoadingScreen", "Niagara", "AsyncMixin", - "ControlFlows" + "ControlFlows", + "StyleTransfer" } );