From 7e227f24f2dcfab7dcc13fe1f29223e0eedec6cf Mon Sep 17 00:00:00 2001
From: Manuel Wagner <m.wagner@aesir-interactive.com>
Date: Wed, 21 Sep 2022 17:02:39 +0200
Subject: [PATCH] added dual style support via new global ShadowMask

---
 Lyra.uproject                                 |   2 +-
 .../Content/NN_PredictorVGG.uasset            |   4 +-
 .../NN_PredictorVGG_SingleStyle.uasset        |   3 +
 .../Content/NN_TransferVGG.uasset             |   4 +-
 .../Content/NN_TransferVGG_SingleStyle.uasset |   3 +
 .../Private/SceneColorToInputTensor.usf       |   2 +-
 .../Private/ShadowMaskToInputTensor.usf       |  29 +++++
 .../StyleTransferSceneViewExtension.cpp       | 116 ++++++++++++++++--
 .../Private/StyleTransferSubsystem.cpp        |  66 +++++++---
 .../Public/StyleTransferSceneViewExtension.h  |   4 +
 .../Public/StyleTransferSubsystem.h           |   7 +-
 .../Private/ShadowMaskToInputTensorCS.cpp     |  19 +++
 .../Public/SceneColorToInputTensorCS.h        |   1 -
 .../Public/ShadowMaskToInputTensorCS.h        |  42 +++++++
 Source/LyraGame/Player/LyraCheatManager.cpp   |  16 +++
 Source/LyraGame/Player/LyraCheatManager.h     |   3 +
 16 files changed, 285 insertions(+), 36 deletions(-)
 create mode 100644 Plugins/StyleTransfer/Content/NN_PredictorVGG_SingleStyle.uasset
 create mode 100644 Plugins/StyleTransfer/Content/NN_TransferVGG_SingleStyle.uasset
 create mode 100644 Plugins/StyleTransfer/Shaders/Private/ShadowMaskToInputTensor.usf
 create mode 100644 Plugins/StyleTransfer/Source/StyleTransferShaders/Private/ShadowMaskToInputTensorCS.cpp
 create mode 100644 Plugins/StyleTransfer/Source/StyleTransferShaders/Public/ShadowMaskToInputTensorCS.h

diff --git a/Lyra.uproject b/Lyra.uproject
index b9e67b7a..c58664b4 100644
--- a/Lyra.uproject
+++ b/Lyra.uproject
@@ -1,6 +1,6 @@
 {
 	"FileVersion": 3,
-	"EngineAssociation": "",
+	"EngineAssociation": "{67760C9E-4E78-132C-8B45-828F7E94A5A9}",
 	"Category": "",
 	"Description": "",
 	"Modules": [
diff --git a/Plugins/StyleTransfer/Content/NN_PredictorVGG.uasset b/Plugins/StyleTransfer/Content/NN_PredictorVGG.uasset
index 10863975..1cd0d566 100644
--- a/Plugins/StyleTransfer/Content/NN_PredictorVGG.uasset
+++ b/Plugins/StyleTransfer/Content/NN_PredictorVGG.uasset
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8e8c479db0bc4d2c91f7a17ba9a82e544068365c77054a15d53ae47a2e52b825
-size 5128046
+oid sha256:b7cb605bc4872bcfc4805f09a3072211029114bd2001a70bb431c2757d4261d4
+size 5128074
diff --git a/Plugins/StyleTransfer/Content/NN_PredictorVGG_SingleStyle.uasset b/Plugins/StyleTransfer/Content/NN_PredictorVGG_SingleStyle.uasset
new file mode 100644
index 00000000..84ea3d55
--- /dev/null
+++ b/Plugins/StyleTransfer/Content/NN_PredictorVGG_SingleStyle.uasset
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3dd9d9ee8f114637fec745b2b1956ae36bf25846fff4255bad04fd3eb9366135
+size 5128122
diff --git a/Plugins/StyleTransfer/Content/NN_TransferVGG.uasset b/Plugins/StyleTransfer/Content/NN_TransferVGG.uasset
index 44b85ead..ecbbba0e 100644
--- a/Plugins/StyleTransfer/Content/NN_TransferVGG.uasset
+++ b/Plugins/StyleTransfer/Content/NN_TransferVGG.uasset
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0ca79b4a94ae300ca98a699ce3ee2e1d6025b418cc83ff6fe48671e564b2b2ee
-size 6797742
+oid sha256:c244fa1e3b61fdd0a8b1ece6e772770bbad766e0a11408734b725811225a6c6b
+size 6888931
diff --git a/Plugins/StyleTransfer/Content/NN_TransferVGG_SingleStyle.uasset b/Plugins/StyleTransfer/Content/NN_TransferVGG_SingleStyle.uasset
new file mode 100644
index 00000000..7de1a355
--- /dev/null
+++ b/Plugins/StyleTransfer/Content/NN_TransferVGG_SingleStyle.uasset
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a97eef14bd1b17fdf5a2ba8aab156a9dfb71790d408b549d23686dc656cd5718
+size 6888979
diff --git a/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf b/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf
index 207d322d..3f1881ab 100644
--- a/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf
+++ b/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
diff --git a/Plugins/StyleTransfer/Shaders/Private/ShadowMaskToInputTensor.usf b/Plugins/StyleTransfer/Shaders/Private/ShadowMaskToInputTensor.usf
new file mode 100644
index 00000000..cf3aa12a
--- /dev/null
+++ b/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"
\ No newline at end of file
diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp
index b2aed579..e3ae5ed9 100644
--- a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp
+++ b/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");
@@ -244,7 +325,7 @@ void FStyleTransferSceneViewExtension::InterpolateTensors(FRDGBuilder& GraphBuil
 		[InterpolateTensorsCS, InterpolateTensorsParameters, InterpolateTensorsThreadGroupCount](FRHICommandList& RHICommandList)
 		{
 			FComputeShaderUtils::Dispatch(RHICommandList, InterpolateTensorsCS,
-										  *InterpolateTensorsParameters, InterpolateTensorsThreadGroupCount);
+			                              *InterpolateTensorsParameters, InterpolateTensorsThreadGroupCount);
 		}
 	);
 }
@@ -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);
 }
diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp
index 1b21b4a1..56ecb378 100644
--- a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp
+++ b/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,12 +91,16 @@ 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());
 	}
-	StyleTransferSceneViewExtension->SetEnabled(true);
+	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;
 
@@ -252,21 +274,29 @@ void UStyleTransferSubsystem::InterpolateStyles(int32 StylePredictionInferenceCo
 			FStyleTransferSceneViewExtension::InterpolateTensors(GraphBuilder, OutputStyleParamsTensor, InputStyleImageTensorA, InputStyleImageTensorB, Alpha);
 		}
 		GraphBuilder.Execute();
-		if(RenderCaptureProvider) RenderCaptureProvider->EndCapture(&RHICommandList);
+		if (RenderCaptureProvider) RenderCaptureProvider->EndCapture(&RHICommandList);
 	});
 }
 
-IRenderCaptureProvider* UStyleTransferSubsystem::ConditionalBeginRenderCapture(FRHICommandListImmediate& RHICommandList)
+IRenderCaptureProvider* BeginRenderCapture(FRHICommandListImmediate& RHICommandList)
+{
+	IRenderCaptureProvider* RenderCaptureProvider = nullptr;
+	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())
 	{
-		const FName RenderCaptureProviderType = IRenderCaptureProvider::GetModularFeatureName();
-		if (IModularFeatures::Get().IsModularFeatureAvailable(RenderCaptureProviderType))
-		{
-			RenderCaptureProvider = &IModularFeatures::Get().GetModularFeature<IRenderCaptureProvider>(RenderCaptureProviderType);
-			RenderCaptureProvider->BeginCapture(&RHICommandList);
-		}
+		RenderCaptureProvider = BeginRenderCapture(RHICommandList);
 	}
 	return RenderCaptureProvider;
 }
diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h
index e03531d8..e90eb7c5 100644
--- a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h
+++ b/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;
 };
diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h
index 4295d443..90db09d5 100644
--- a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h
+++ b/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);
\ No newline at end of file
diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/ShadowMaskToInputTensorCS.cpp b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/ShadowMaskToInputTensorCS.cpp
new file mode 100644
index 00000000..bac0a554
--- /dev/null
+++ b/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
diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/SceneColorToInputTensorCS.h b/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/SceneColorToInputTensorCS.h
index 616e7851..ab555686 100644
--- a/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/SceneColorToInputTensorCS.h
+++ b/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)
diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/ShadowMaskToInputTensorCS.h b/Plugins/StyleTransfer/Source/StyleTransferShaders/Public/ShadowMaskToInputTensorCS.h
new file mode 100644
index 00000000..96cd1b09
--- /dev/null
+++ b/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:
+};
diff --git a/Source/LyraGame/Player/LyraCheatManager.cpp b/Source/LyraGame/Player/LyraCheatManager.cpp
index 329316f9..0b60d47c 100644
--- a/Source/LyraGame/Player/LyraCheatManager.cpp
+++ b/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);
+}
diff --git a/Source/LyraGame/Player/LyraCheatManager.h b/Source/LyraGame/Player/LyraCheatManager.h
index f193078e..c00229ec 100644
--- a/Source/LyraGame/Player/LyraCheatManager.h
+++ b/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;