added dual style support via new global ShadowMask

Lyra.uproject

@@ -1,6 +1,6 @@
 {
 	"FileVersion": 3,
-	"EngineAssociation": "",
+	"EngineAssociation": "{67760C9E-4E78-132C-8B45-828F7E94A5A9}",
 	"Category": "",
 	"Description": "",
 	"Modules": [

Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf

@@ -1,6 +1,6 @@
 // Copyright 2022 Manuel Wagner - All rights reserved
 
-Texture2D<float4> InputTexture;
+Texture2D InputTexture;
 SamplerState InputTextureSampler;
 RWBuffer<float> OutputUAV;
 uint2 OutputDimensions; // X = InputTensor.GetSize(1), Y = InputTensor.GetSize(2) -> this does not correspond to input texture XY

Plugins/StyleTransfer/Shaders/Private/ShadowMaskToInputTensor.usf

@@ -0,0 +1,29 @@
+// Copyright 2022 Manuel Wagner - All rights reserved
+
+Texture2D InputTexture;
+SamplerState InputTextureSampler;
+RWBuffer<float> OutputUAV;
+uint2 OutputDimensions; // X = InputTensor.GetSize(1), Y = InputTensor.GetSize(2) -> this does not correspond to input texture XY
+float2 HalfPixelUV;
+
+[numthreads(THREADGROUP_SIZE_X, THREADGROUP_SIZE_Y, THREADGROUP_SIZE_Z)]
+void ShadowMaskToInputTensorCS(in const uint3 DispatchThreadID : SV_DispatchThreadID)
+{
+	const uint2 OutputUAVTexelCoordinate = DispatchThreadID.xy;
+	if(any(OutputUAVTexelCoordinate >= OutputDimensions))
+	{
+		return;
+	}
+
+	const uint GlobalIndex = ((OutputUAVTexelCoordinate.x * OutputDimensions.y + OutputUAVTexelCoordinate.y));
+
+	// note that the OutputUAV has shape (1, Y, X, C)
+	// which is why we need to flip the indexing
+	const float2 UV = float2(OutputUAVTexelCoordinate.yx) / float2(OutputDimensions.yx) + HalfPixelUV;
+
+	const float4 TextureValue = InputTexture.SampleLevel(InputTextureSampler, UV, 0);
+
+	OutputUAV[GlobalIndex + 0] = TextureValue.r;
+}
+
+#include "/Engine/Public/Platform.ush"

Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp

@@ -23,8 +23,17 @@
 #include "PostProcess/PostProcessMaterial.h"
 #include "OutputTensorToSceneColorCS.h"
 #include "PixelShaderUtils.h"
+#include "RendererUtils.h"
 #include "SceneColorToInputTensorCS.h"
+#include "ShadowMaskToInputTensorCS.h"
 #include "StyleTransferModule.h"
+#include "StyleTransferSubsystem.h"
+
+TAutoConsoleVariable<bool> CVarAutoCaptureStyleTransfer(
+	TEXT("r.StyleTransfer.AutoCaptureTransfer"),
+	false,
+	TEXT("Set to true to automatically capture the style transfer when it is done")
+);
 
 template <class OutType, class InType>
 OutType CastNarrowingSafe(InType InValue)
@@ -49,6 +58,17 @@ FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRe
 	  , InferenceContext(InInferenceContext)
 {
 	ensure(InStyleTransferNetwork->GetDeviceType() == ENeuralDeviceType::GPU);
+
+	for (uint32 i = 0; i < InStyleTransferNetwork->GetInputTensorNumber(); i++)
+	{
+		const FString& TensorName = InStyleTransferNetwork->GetInputTensor(i).GetName();
+		if (TensorName == "content") ContentInputTensorIndex = i;
+		else if (TensorName == "style_weights") StyleWeightsInputTensorIndex = i;
+		else if (TensorName == "style_params") StyleParamsInputTensorIndex = i;
+	}
+	check(ContentInputTensorIndex != INDEX_NONE);
+	check(StyleWeightsInputTensorIndex != INDEX_NONE);
+	check(StyleParamsInputTensorIndex != INDEX_NONE);
 }
 
 bool FStyleTransferSceneViewExtension::IsActiveThisFrame_Internal(const FSceneViewExtensionContext& Context) const
@@ -181,7 +201,54 @@ FRDGTexture* FStyleTransferSceneViewExtension::TensorToTexture(FRDGBuilder& Grap
 	return OutputTexture;
 }
 
-void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor)
+FRDGTexture* TensorToTexture(FRDGBuilder& GraphBuilder, const FRDGTextureDesc& BaseDestinationDesc, const FNeuralTensor& SourceTensor)
+{
+	FIntVector SourceTensorDimensions = {
+		CastNarrowingSafe<int32>(SourceTensor.GetSize(1)),
+		CastNarrowingSafe<int32>(SourceTensor.GetSize(2)),
+		CastNarrowingSafe<int32>(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<FOutputTensorToSceneColorCS::FParameters>();
+	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<FOutputTensorToSceneColorCS> 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;
+}
+
+template <class ShaderType = FSceneColorToInputTensorCS>
+void TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor)
 {
 	const FIntVector InputTensorDimensions = {
 		CastNarrowingSafe<int32>(DestinationTensor.GetSize(1)),
@@ -191,7 +258,7 @@ void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder
 	const FIntPoint SceneColorRenderTargetDimensions = SourceTexture->Desc.Extent;
 
 	FRDGBufferRef StyleTransferContentInputBuffer = GraphBuilder.RegisterExternalBuffer(DestinationTensor.GetPooledBuffer());
-	auto SceneColorToInputTensorParameters = GraphBuilder.AllocParameters<FSceneColorToInputTensorCS::FParameters>();
+	typename ShaderType::FParameters* SceneColorToInputTensorParameters = GraphBuilder.AllocParameters<typename ShaderType::FParameters>();
 	SceneColorToInputTensorParameters->TensorVolume = CastNarrowingSafe<uint32>(DestinationTensor.Num());
 	SceneColorToInputTensorParameters->InputTexture = SourceTexture;
 	SceneColorToInputTensorParameters->InputTextureSampler = TStaticSamplerState<SF_Bilinear>::GetRHI();
@@ -200,12 +267,13 @@ void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder
 	SceneColorToInputTensorParameters->HalfPixelUV = FVector2f(0.5f / SceneColorRenderTargetDimensions.X, 0.5 / SceneColorRenderTargetDimensions.Y);
 	FIntVector SceneColorToInputTensorGroupCount = FComputeShaderUtils::GetGroupCount(
 		{InputTensorDimensions.X, InputTensorDimensions.Y, 1},
-		FSceneColorToInputTensorCS::ThreadGroupSize
+		ShaderType::ThreadGroupSize
 	);
 
-	TShaderMapRef<FSceneColorToInputTensorCS> SceneColorToInputTensorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel));
+	TShaderMapRef<ShaderType> SceneColorToInputTensorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel));
+
 	GraphBuilder.AddPass(
-		RDG_EVENT_NAME("TextureToTensor"),
+		RDG_EVENT_NAME("TextureToTensor(%s)", ShaderType::StaticType.GetName()),
 		SceneColorToInputTensorParameters,
 		ERDGPassFlags::Compute,
 		[SceneColorToInputTensorCS, SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount](FRHICommandList& RHICommandList)
@@ -216,6 +284,19 @@ void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder
 	);
 }
 
+template <uint8 Channels>
+void TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor);
+template <>
+void TextureToTensor<1>(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor){return TextureToTensor<FShadowMaskToInputTensorCS>(GraphBuilder, SourceTexture, DestinationTensor);}
+template <>
+void TextureToTensor<3>(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor){return TextureToTensor<FSceneColorToInputTensorCS>(GraphBuilder, SourceTexture, DestinationTensor);}
+
+
+void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor)
+{
+	::TextureToTensor<3>(GraphBuilder, SourceTexture, DestinationTensor);
+}
+
 void FStyleTransferSceneViewExtension::InterpolateTensors(FRDGBuilder& GraphBuilder, const FNeuralTensor& DestinationTensor, const FNeuralTensor& InputTensorA, const FNeuralTensor& InputTensorB, float Alpha)
 {
 	RDG_EVENT_SCOPE(GraphBuilder, "InterpolateTensors");
@@ -265,12 +346,23 @@ FScreenPassTexture FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap
 
 	RDG_EVENT_SCOPE(GraphBuilder, "StyleTransfer");
 
+	IRenderCaptureProvider* RenderCaptureProvider = nullptr;
+	if (CVarAutoCaptureStyleTransfer.GetValueOnRenderThread())
+	{
+		RenderCaptureProvider = BeginRenderCapture(GraphBuilder.RHICmdList);
+	}
+
 	checkSlow(View.bIsViewInfo);
 	const FViewInfo& ViewInfo = static_cast<const FViewInfo&>(View);
 
-	const FNeuralTensor& StyleTransferContentInputTensor = StyleTransferNetwork->GetInputTensorForContext(*InferenceContext, 0);
+	const FNeuralTensor& StyleTransferContentInputTensor = StyleTransferNetwork->GetInputTensorForContext(*InferenceContext, ContentInputTensorIndex);
+	const FNeuralTensor& StyleTransferStyleWeightsInputTensor = StyleTransferNetwork->GetInputTensorForContext(*InferenceContext, StyleWeightsInputTensorIndex);
 
-	TextureToTensor(GraphBuilder, SceneColor.Texture, StyleTransferContentInputTensor);
+	check(GScreenShadowMaskTexture);
+	::TextureToTensor<1>(GraphBuilder, GScreenShadowMaskTexture, StyleTransferStyleWeightsInputTensor);
+	GScreenShadowMaskTexture = nullptr;
+
+	::TextureToTensor<3>(GraphBuilder, SceneColor.Texture, StyleTransferContentInputTensor);
 
 	StyleTransferNetwork->Run(GraphBuilder, *InferenceContext);
 
@@ -294,5 +386,13 @@ FScreenPassTexture FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap
 		BackBufferRenderTarget = StyleTransferOutputTarget;
 	}
 
+	if (RenderCaptureProvider)
+	{
+		/*GraphBuilder.AddPass(RDG_EVENT_NAME("EndCapture"), ERDGPassFlags::None, [RenderCaptureProvider](FRHICommandListImmediate& RHICommandList)
+		{
+			RenderCaptureProvider->EndCapture(&RHICommandList);
+		});*/
+	}
+
 	return MoveTemp(*BackBufferRenderTarget);
 }

Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp

@@ -42,6 +42,7 @@ void UStyleTransferSubsystem::Deinitialize()
 
 bool UStyleTransferSubsystem::Tick(float DeltaTime)
 {
+	return false;
 	if (!GetWorld())
 		return true;
 
@@ -74,12 +75,14 @@ void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportCl
 
 		if (!StyleTransferInferenceContext || *StyleTransferInferenceContext == INDEX_NONE)
 		{
+			UE_LOG(LogStyleTransfer, Log, TEXT("Creating Inference Context for StyleTransfer"));
 			StyleTransferInferenceContext = MakeShared<int32>(StyleTransferNetwork->CreateInferenceContext());
 			checkf(*StyleTransferInferenceContext != INDEX_NONE, TEXT("Could not create inference context for StyleTransferNetwork"));
 		}
 
-		for (int32 i = 0; i < FMath::Min(2, StyleTransferSettings->StyleTextures.Num()); ++i)
+		for (uint32 i = 0; i < FMath::Min(2u, uint32(StyleTransferSettings->StyleTextures.Num())); ++i)
 		{
+			UE_LOG(LogStyleTransfer, Log, TEXT("Creating Inference Context for Style %i"), i);
 			const int32& StylePredictionInferenceContext = StylePredictionInferenceContexts.Emplace_GetRef(StylePredictionNetwork->CreateInferenceContext());
 			checkf(StylePredictionInferenceContext != INDEX_NONE, TEXT("Could not create inference context for StylePredictionNetwork"));
 
@@ -88,13 +91,17 @@ void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportCl
 #if WITH_EDITOR
 			FTextureCompilingManager::Get().FinishCompilation({StyleTexture});
 #endif
-			UpdateStyle(StyleTexture, StylePredictionInferenceContext);
+			UpdateStyle(StyleTexture, i, StylePredictionInferenceContext);
 		}
 		//UpdateStyle(FPaths::GetPath("C:\\projects\\realtime-style-transfer\\temp\\style_params_tensor.bin"));
+		UE_LOG(LogStyleTransfer, Log, TEXT("Creating FStyleTransferSceneViewExtension"));
 		StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension<FStyleTransferSceneViewExtension>(ViewportClient->GetWorld(), ViewportClient, StyleTransferNetwork, StyleTransferInferenceContext.ToSharedRef());
 	}
+	if (StyleTransferSceneViewExtension)
+	{
 		StyleTransferSceneViewExtension->SetEnabled(true);
 	}
+}
 
 void UStyleTransferSubsystem::StopStylizingViewport()
 {
@@ -116,12 +123,17 @@ void UStyleTransferSubsystem::StopStylizingViewport()
 	}
 }
 
-void UStyleTransferSubsystem::UpdateStyle(UTexture2D* StyleTexture, int32 StylePredictionInferenceContext)
+BEGIN_SHADER_PARAMETER_STRUCT(FCopyBufferParameters,)
+	RDG_BUFFER_ACCESS(SrcBuffer, ERHIAccess::CopySrc)
+	RDG_BUFFER_ACCESS(DstBuffer, ERHIAccess::CopyDest)
+END_SHADER_PARAMETER_STRUCT()
+
+void UStyleTransferSubsystem::UpdateStyle(UTexture2D* StyleTexture, uint32 StyleIndex, int32 StylePredictionInferenceContext)
 {
 	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, StylePredictionInferenceContext](FRHICommandListImmediate& RHICommandList)
+	ENQUEUE_RENDER_COMMAND(StylePrediction)([this, StyleTexture, StylePredictionInferenceContext, StyleIndex](FRHICommandListImmediate& RHICommandList)
 	{
 		IRenderCaptureProvider* RenderCaptureProvider = ConditionalBeginRenderCapture(RHICommandList);
 
@@ -142,11 +154,21 @@ void UStyleTransferSubsystem::UpdateStyle(UTexture2D* StyleTexture, int32 StyleP
 			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 uint64 DstOffset = StyleIndex * NumBytes;
 
-			AddCopyBufferPass(GraphBuilder, InputStyleParamsBuffer, OutputStyleParamsBuffer);
+
+			FCopyBufferParameters* Parameters = GraphBuilder.AllocParameters<FCopyBufferParameters>();
+			Parameters->SrcBuffer = OutputStyleParamsBuffer;
+			Parameters->DstBuffer = InputStyleParamsBuffer;
+
+			GraphBuilder.AddPass(
+				RDG_EVENT_NAME("CopyBuffer(%s Size=%ubytes)", Parameters->SrcBuffer, Parameters->SrcBuffer->Desc.GetSize()),
+				Parameters,
+				ERDGPassFlags::Copy,
+				[Parameters, NumBytes, DstOffset](FRHICommandList& RHICmdList)
+				{
+					RHICmdList.CopyBufferRegion(Parameters->DstBuffer->GetRHI(), DstOffset, Parameters->SrcBuffer->GetRHI(), 0, NumBytes);
+				});
 		}
 		GraphBuilder.Execute();
 
@@ -163,7 +185,7 @@ void UStyleTransferSubsystem::UpdateStyle(FString StyleTensorDataPath)
 {
 	FArchive& FileReader = *IFileManager::Get().CreateFileReader(*StyleTensorDataPath);
 	TArray<float> StyleParams;
-	StyleParams.SetNumUninitialized(192);
+	StyleParams.SetNumUninitialized(2758);
 
 	FileReader << StyleParams;
 
@@ -256,17 +278,25 @@ void UStyleTransferSubsystem::InterpolateStyles(int32 StylePredictionInferenceCo
 	});
 }
 
-IRenderCaptureProvider* UStyleTransferSubsystem::ConditionalBeginRenderCapture(FRHICommandListImmediate& RHICommandList)
+IRenderCaptureProvider* BeginRenderCapture(FRHICommandListImmediate& RHICommandList)
 {
 	IRenderCaptureProvider* RenderCaptureProvider = nullptr;
-	if (CVarAutoCaptureStylePrediction.GetValueOnRenderThread())
-	{
 	const FName RenderCaptureProviderType = IRenderCaptureProvider::GetModularFeatureName();
 	if (IModularFeatures::Get().IsModularFeatureAvailable(RenderCaptureProviderType))
 	{
 		RenderCaptureProvider = &IModularFeatures::Get().GetModularFeature<IRenderCaptureProvider>(RenderCaptureProviderType);
+		RenderCaptureProvider->EndCapture(&RHICommandList);
 		RenderCaptureProvider->BeginCapture(&RHICommandList);
 	}
-	}
+	return RenderCaptureProvider;
+}
+
+IRenderCaptureProvider* ConditionalBeginRenderCapture(FRHICommandListImmediate& RHICommandList)
+{
+	IRenderCaptureProvider* RenderCaptureProvider = nullptr;
+	if (CVarAutoCaptureStylePrediction.GetValueOnRenderThread())
+	{
+		RenderCaptureProvider = BeginRenderCapture(RHICommandList);
+	}
 	return RenderCaptureProvider;
 }

Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h

@@ -57,4 +57,8 @@ private:
 	bool bIsEnabled = true;
 
 	int32 NumFramesCaptured = -1;
+
+	int32 ContentInputTensorIndex = INDEX_NONE;
+	int32 StyleWeightsInputTensorIndex = INDEX_NONE;
+	int32 StyleParamsInputTensorIndex = INDEX_NONE;
 };

Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h

@@ -30,7 +30,7 @@ public:
 	void StartStylizingViewport(FViewportClient* ViewportClient);
 	void StopStylizingViewport();
 
-	void UpdateStyle(UTexture2D* StyleTexture, int32 StylePredictionInferenceContext);
+	void UpdateStyle(UTexture2D* StyleTexture, uint32 StyleIndex, int32 StylePredictionInferenceContext);
 	void UpdateStyle(FString StyleTensorDataPath);
 	void InterpolateStyles(int32 StylePredictionInferenceContextA, int32 StylePredictionInferenceContextB, float Alpha);
 
@@ -52,6 +52,7 @@ private:
 	void HandleConsoleVariableChanged(IConsoleVariable*);
 
 	void LoadNetworks();
-
-	IRenderCaptureProvider* ConditionalBeginRenderCapture(FRHICommandListImmediate& RHICommandList);
 };
+
+IRenderCaptureProvider* BeginRenderCapture(FRHICommandListImmediate& RHICommandList);
+IRenderCaptureProvider* ConditionalBeginRenderCapture(FRHICommandListImmediate& RHICommandList);

Plugins/StyleTransfer/Source/StyleTransferShaders/Private/ShadowMaskToInputTensorCS.cpp

@@ -0,0 +1,19 @@
+// Copyright Manuel Wagner All Rights Reserved.
+
+#include "ShadowMaskToInputTensorCS.h"
+
+const FIntVector FShadowMaskToInputTensorCS::ThreadGroupSize{8, 8, 1};
+
+void FShadowMaskToInputTensorCS::ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment)
+{
+	FGlobalShader::ModifyCompilationEnvironment(Parameters, OutEnvironment);
+
+	OutEnvironment.SetDefine(TEXT("THREADGROUP_SIZE_X"), ThreadGroupSize.X);
+	OutEnvironment.SetDefine(TEXT("THREADGROUP_SIZE_Y"), ThreadGroupSize.Y);
+	OutEnvironment.SetDefine(TEXT("THREADGROUP_SIZE_Z"), ThreadGroupSize.Z);
+}
+
+
+IMPLEMENT_GLOBAL_SHADER(FShadowMaskToInputTensorCS,
+						"/Plugins/StyleTransfer/Shaders/Private/ShadowMaskToInputTensor.usf",
+						"ShadowMaskToInputTensorCS", SF_Compute); // Path defined in StyleTransferModule.cpp

Plugins/StyleTransfer/Source/StyleTransferShaders/Public/SceneColorToInputTensorCS.h

@@ -19,7 +19,6 @@ class STYLETRANSFERSHADERS_API FSceneColorToInputTensorCS : public FGlobalShader
 
 	static const FIntVector ThreadGroupSize;
 
-
 	BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
 		// Input variables
 		SHADER_PARAMETER(uint32, TensorVolume)

Plugins/StyleTransfer/Source/StyleTransferShaders/Public/ShadowMaskToInputTensorCS.h

@@ -0,0 +1,42 @@
+// Copyright Manuel Wagner 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 FShadowMaskToInputTensorCS : public FGlobalShader
+{
+	DECLARE_GLOBAL_SHADER(FShadowMaskToInputTensorCS);
+	SHADER_USE_PARAMETER_STRUCT(FShadowMaskToInputTensorCS, FGlobalShader)
+
+
+	static const FIntVector ThreadGroupSize;
+
+	BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
+		// Input variables
+		SHADER_PARAMETER(uint32, TensorVolume)
+		// SRV/UAV variables
+		SHADER_PARAMETER_RDG_BUFFER_UAV(RWBuffer<float>, OutputUAV)
+		// Optional SRV/UAV variables
+		SHADER_PARAMETER_RDG_TEXTURE(Texture2D, InputTexture)
+		SHADER_PARAMETER_SAMPLER(SamplerState, InputTextureSampler)
+		SHADER_PARAMETER(FIntPoint, OutputDimensions)
+		SHADER_PARAMETER(FVector2f, HalfPixelUV)
+	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:
+};

Source/LyraGame/Player/LyraCheatManager.cpp

@@ -14,6 +14,7 @@
 #include "AbilitySystem/LyraAbilitySystemComponent.h"
 #include "AbilitySystemGlobals.h"
 #include "GameplayEffect.h"
+#include "HighResScreenshot.h"
 #include "Character/LyraHealthComponent.h"
 #include "Character/LyraPawnExtensionComponent.h"
 #include "System/LyraSystemStatics.h"
@@ -422,3 +423,18 @@ void ULyraCheatManager::UnlimitedHealth(int32 Enabled)
 		}
 	}
 }
+
+
+void ULyraCheatManager::Singinwhale_StartScreenshotInterval(float IntervalSeconds)
+{
+	FTimerHandle _Handle = FTimerHandle();
+	GetWorld()->GetTimerManager().SetTimer(_Handle, FTimerDelegate::CreateWeakLambda(this, [this]()
+	{
+		FHighResScreenshotConfig& HighResScreenshotConfig = GetHighResScreenshotConfig();
+		HighResScreenshotConfig.bDumpBufferVisualizationTargets = true;
+		HighResScreenshotConfig.bDateTimeBasedNaming = true;
+		HighResScreenshotConfig.SetResolution(1920,1080);
+		HighResScreenshotConfig.SetHDRCapture(true);
+		FScreenshotRequest::RequestScreenshot(true);
+	}), IntervalSeconds, true, 0);
+}

Source/LyraGame/Player/LyraCheatManager.h

@@ -95,6 +95,9 @@ public:
 	UFUNCTION(Exec, BlueprintAuthorityOnly)
 	virtual void UnlimitedHealth(int32 Enabled = -1);
 
+	UFUNCTION(Exec)
+	void Singinwhale_StartScreenshotInterval(float IntervalSeconds);
+
 protected:
 
 	virtual void EnableDebugCamera() override;