made style transfer world specific

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

@@ -40,8 +40,8 @@ OutType CastNarrowingSafe(InType InValue)
 }
 
 
-FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef<int32> InInferenceContext)
+FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, UWorld* World, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef<int32> InInferenceContext)
-	: FSceneViewExtensionBase(AutoRegister)
+	: FWorldSceneViewExtension(AutoRegister, World)
 	  , StyleTransferNetworkWeakPtr(InStyleTransferNetwork)
 	  , StyleTransferNetwork(InStyleTransferNetwork)
 	  , LinkedViewportClient(AssociatedViewportClient)
@@ -50,14 +50,12 @@ FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRe
 	ensure(InStyleTransferNetwork->GetDeviceType() == ENeuralDeviceType::GPU);
 }
 
-void FStyleTransferSceneViewExtension::SetupViewFamily(FSceneViewFamily& InViewFamily)
-{
-}
-
 bool FStyleTransferSceneViewExtension::IsActiveThisFrame_Internal(const FSceneViewExtensionContext& Context) const
 {
 	check(IsInGameThread());
-	return bIsEnabled && *InferenceContext != -1 && StyleTransferNetworkWeakPtr.IsValid();
+	return FWorldSceneViewExtension::IsActiveThisFrame_Internal(Context)
+		&& bIsEnabled
+		&& *InferenceContext != -1 && StyleTransferNetworkWeakPtr.IsValid();
 }
 
 void FStyleTransferSceneViewExtension::AddRescalingTextureCopy(FRDGBuilder& GraphBuilder, FRDGTexture& RDGSourceTexture, FScreenPassRenderTarget& DestinationRenderTarget)
@@ -115,22 +113,22 @@ void FStyleTransferSceneViewExtension::PreRenderViewFamily_RenderThread(FRDGBuil
 {
 	return;
 	const FName RenderCaptureProviderType = IRenderCaptureProvider::GetModularFeatureName();
-	if(!IModularFeatures::Get().IsModularFeatureAvailable(RenderCaptureProviderType))
+	if (!IModularFeatures::Get().IsModularFeatureAvailable(RenderCaptureProviderType))
 		return;
 
 	IRenderCaptureProvider& RenderCaptureProvider = IModularFeatures::Get().GetModularFeature<IRenderCaptureProvider>(RenderCaptureProviderType);
-	if(bIsEnabled && NumFramesCaptured == -1)
+	if (bIsEnabled && NumFramesCaptured == -1)
 	{
 		RenderCaptureProvider.BeginCapture(&GRHICommandList.GetImmediateCommandList());
 		NumFramesCaptured = 0;
 	}
 
-	if(NumFramesCaptured >= 0)
+	if (NumFramesCaptured >= 0)
 	{
 		++NumFramesCaptured;
 	}
 
-	if(NumFramesCaptured == 10)
+	if (NumFramesCaptured == 10)
 	{
 		RenderCaptureProvider.EndCapture(&GRHICommandList.GetImmediateCommandList());
 	}
@@ -219,8 +217,6 @@ void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder
 
 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))

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

@@ -52,10 +52,12 @@ void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportCl
 
 		UTexture2D* StyleTexture = StyleTransferSettings->StyleTexture.LoadSynchronous();
 		//UTexture2D* StyleTexture = LoadObject<UTexture2D>(this, TEXT("/Script/Engine.Texture2D'/StyleTransfer/T_StyleImage.T_StyleImage'"));
+#if WITH_EDITOR
 		FTextureCompilingManager::Get().FinishCompilation({StyleTexture});
+#endif
 		UpdateStyle(StyleTexture);
 		//UpdateStyle(FPaths::GetPath("C:\\projects\\realtime-style-transfer\\temp\\style_params_tensor.bin"));
-		StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension<FStyleTransferSceneViewExtension>(ViewportClient, StyleTransferNetwork, StyleTransferInferenceContext.ToSharedRef());
+		StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension<FStyleTransferSceneViewExtension>(ViewportClient->GetWorld(), ViewportClient, StyleTransferNetwork, StyleTransferInferenceContext.ToSharedRef());
 
 	}
 	StyleTransferSceneViewExtension->SetEnabled(true);
@@ -68,6 +70,7 @@ void UStyleTransferSubsystem::StopStylizingViewport()
 	if (StylePredictionInferenceContext != INDEX_NONE)
 	{
 		StylePredictionNetwork->DestroyInferenceContext(StylePredictionInferenceContext);
+		StylePredictionInferenceContext = INDEX_NONE;
 	}
 	if (StyleTransferInferenceContext && *StyleTransferInferenceContext != INDEX_NONE)
 	{

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

@@ -5,20 +5,25 @@ struct FNeuralTensor;
 struct FScreenPassRenderTarget;
 class UNeuralNetwork;
 
-class FStyleTransferSceneViewExtension : public FSceneViewExtensionBase
+class FStyleTransferSceneViewExtension : public FWorldSceneViewExtension
 {
 public:
 	using Ptr = TSharedPtr<FStyleTransferSceneViewExtension, ESPMode::ThreadSafe>;
 	using Ref = TSharedRef<FStyleTransferSceneViewExtension, ESPMode::ThreadSafe>;
 
-	FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef<int32> InInferenceContext);
+	FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, UWorld* World, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef<int32> InInferenceContext);
 
 	// - ISceneViewExtension
 	virtual void SubscribeToPostProcessingPass(EPostProcessingPass Pass, FAfterPassCallbackDelegateArray& InOutPassCallbacks, bool bIsPassEnabled) override;
+
 	virtual void PreRenderViewFamily_RenderThread(FRDGBuilder& GraphBuilder, FSceneViewFamily& InViewFamily) override;
+
 	FScreenPassTexture PostProcessPassAfterTonemap_RenderThread(FRDGBuilder& GraphBuilder, const FSceneView& View,
 	                                                            const FPostProcessMaterialInputs& InOutInputs);
-	virtual void SetupViewFamily(FSceneViewFamily& InViewFamily) override;
+
+	virtual void SetupViewFamily(FSceneViewFamily& InViewFamily) override
+	{
+	}
 
 	virtual void SetupView(FSceneViewFamily& InViewFamily, FSceneView& InView) override
 	{
@@ -35,7 +40,6 @@ public:
 	bool IsEnabled() const { return bIsEnabled; }
 
 
-
 	static void AddRescalingTextureCopy(FRDGBuilder& GraphBuilder, FRDGTexture& RDGSourceTexture, FScreenPassRenderTarget& DestinationRenderTarget);
 	static FRDGTexture* TensorToTexture(FRDGBuilder& GraphBuilder, const FRDGTextureDesc& BaseDestinationDesc, const FNeuralTensor& SourceTensor);
 	static void TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor);

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

@@ -1,7 +1,6 @@
 // Copyright Epic Games, Inc. All Rights Reserved.
 
 #include "OutputTensorToSceneColorCS.h"
-#include "Utils.h"
 
 const FIntVector FOutputTensorToSceneColorCS::ThreadGroupSize{8, 8, 1};
 

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

@@ -1,7 +1,6 @@
 // Copyright Epic Games, Inc. All Rights Reserved.
 
 #include "SceneColorToInputTensorCS.h"
-#include "Utils.h"
 
 const FIntVector FSceneColorToInputTensorCS::ThreadGroupSize{8, 8, 1};