Made dependency installation cancellable and fixed some issues with clearing system dependencies

Author: Mystfit

StableDiffusionTools/Content/Python/bridges/DiffusersBridge.py

@@ -224,7 +224,10 @@ def InitModel(self, new_model_options, allow_nsfw, padding_mode):
         # Performance options for low VRAM gpus
         #self.pipe.enable_sequential_cpu_offload()
         self.pipe.enable_attention_slicing()
-        self.pipe.unet = torch.compile(self.pipe.unet)
+        try:
+            self.pipe.unet = torch.compile(self.pipe.unet)
+        except RuntimeError as e:
+            print(f"WARNING: Couldn't compile unet for model. Exception given was '{e}'")
         #self.pipe.enable_xformers_memory_efficient_attention()
         if hasattr(self.pipe, "enable_model_cpu_offload"):
             self.pipe.enable_model_cpu_offload()
@@ -248,22 +251,14 @@ def InitModel(self, new_model_options, allow_nsfw, padding_mode):
         return result
 
     def InitUpsampler(self):
-        #upsampler = None
-        #try:
-        #    upsampler = RealESRGANModel.from_pretrained("nateraw/real-esrgan")
-        #    upsampler = upsampler.to("cuda")
-        #except Exception as e:
-        #    print("Could not load upsampler. Exception was ".format(e))
-        #print(upsampler)
-        #return upsampler
-        upsampler = StableDiffusionUpscalePipeline.from_pretrained(
-            "stabilityai/stable-diffusion-x4-upscaler", revision="fp16", torch_dtype=torch.float16
-        )
-        upsampler.enable_sequential_cpu_offload()
-        upsampler.enable_attention_slicing()
+        upsampler = None
+        try:
+            upsampler = RealESRGANModel.from_pretrained("nateraw/real-esrgan")
+            upsampler = upsampler.to("cuda")
+        except Exception as e:
+            print("Could not load upsampler. Exception was ".format(e))
         return upsampler
 
-
     @unreal.ufunction(override=True)
     def ReleaseModel(self):
         if hasattr(self, "pipe"):
@@ -331,17 +326,11 @@ def GenerateImageFromStartImage(self, input):
         max_seed = abs(int((2**31) / 2) - 1)
         seed = random.randrange(0, max_seed) if input.options.seed < 0 else input.options.seed
 
-        # Collate prompts
-        positive_prompts = " ".join([f"({split_p.strip()}){prompt.weight}" for prompt in input.options.positive_prompts for split_p in prompt.prompt.split(",")])
-        negative_prompts = " ".join([f"({split_p.strip()}){prompt.weight}" for prompt in input.options.negative_prompts for split_p in prompt.prompt.split(",")])
-        print(positive_prompts)
-        print(negative_prompts)
-        positive_prompt_tensors = self.compel.build_conditioning_tensor(positive_prompts if positive_prompts else "")
-        negative_prompt_tensors = self.compel.build_conditioning_tensor(negative_prompts if negative_prompts else "")
-        prompt_tensors = torch.cat(self.compel.pad_conditioning_tensors_to_same_length([positive_prompt_tensors, negative_prompt_tensors]))
+        # Create prompt
+        prompt_tensors = self.build_prompt_tensors(positive_prompts=input.options.positive_prompts, negative_prompts=input.options.negative_prompts, compel=self.compel)
 
         with torch.inference_mode():
-            with autocast("cuda", dtype=torch.float16):
+            #with autocast("cuda", dtype=torch.float16):
                 generator = torch.Generator(device="cpu")
                 generator.manual_seed(seed)
                 generation_args = {
@@ -422,6 +411,16 @@ def ImageProgressStep(self, step: int, timestep: int, latents: torch.FloatTensor
         if self.abort and self.executor:
             self.executor.stop()
 
+    def build_prompt_tensors(self, positive_prompts: list[unreal.Prompt], negative_prompts: list[unreal.Prompt], compel: Compel):
+         # Collate prompts
+        positive_prompts = " ".join([f"({split_p.strip()}){prompt.weight}" for prompt in positive_prompts for split_p in prompt.prompt.split(",")])
+        negative_prompts = " ".join([f"({split_p.strip()}){prompt.weight}" for prompt in negative_prompts for split_p in prompt.prompt.split(",")])
+        print(positive_prompts)
+        print(negative_prompts)
+        positive_prompt_tensors = compel.build_conditioning_tensor(positive_prompts if positive_prompts else "")
+        negative_prompt_tensors = compel.build_conditioning_tensor(negative_prompts if negative_prompts else "")
+        return torch.cat(compel.pad_conditioning_tensors_to_same_length([positive_prompt_tensors, negative_prompt_tensors]))
+
     @unreal.ufunction(override=True)
     def StopImageGeneration(self):
         self.abort = True

StableDiffusionTools/Content/Python/bridges/DiffusersBridge_dependencies.py

@@ -17,5 +17,6 @@ def MakeDependency(options):
         MakeDependency({"name":"accelerate", "version": "0.18.0"}),
         MakeDependency({"name":"diffusers", "version": "0.15.0"}),
         MakeDependency({"name":"compel", "version": "1.1.3"}),
-        MakeDependency({"name":"opencv-python", "version":"4.7.0.72"})
+        MakeDependency({"name":"opencv-python", "version":"4.7.0.72"}),
+        MakeDependency({"name":"realesrgan", "version": "0.3.0"})
     ]

StableDiffusionTools/Content/Python/dependency_manager.py

@@ -126,17 +126,17 @@ def all_dependencies_installed(self):
         print(f"All dependencies installed? {dependencies_installed}")
         return dependencies_installed
 
-    def clear_all_dependencies(self, env_dir):
+    def clear_all_dependencies(self, env_dir, reset_system_deps):
         # Remove external site-packages dir
         if os.path.exists(env_dir):
             shutil.rmtree(env_dir)
 
         # Remove any leftover packages that are still in Unreal's base site-packages dir. 
-        # Only valid if we're upgrading our plugin from any version before 0.8.2
-        if parse_version(self.get_plugin_version_name()) >= parse_version('0.8.2'):
+        # Only really valid if we're upgrading our plugin from any version before 0.8.2
+        if reset_system_deps:
             with open(Path(os.path.dirname(__file__)) / "requirements.txt") as f:
-                requirements = f.read().splitlines()
-                cmd = [f"{pythonpath}", '-m', 'pip', 'uninstall'] + requirements
+                requirements = [package.split('==')[0] for package in f.read().splitlines()]
+                cmd = [f"{pythonpath}", '-m', 'pip', 'uninstall'] + requirements + ['-y']
                 cmd_string = " ".join(cmd)
                 print(f"Uninstalling dependencies using command '{cmd_string}'")
                 try: 

StableDiffusionTools/Content/Python/init_unreal.py

@@ -79,9 +79,11 @@ def load_manifests(manifest_dir):
 
 # Nuke dependencies before loading them if we're trying to reset the editor dependencies
 reset_deps = dependency_options.get_editor_property("ClearDependenciesOnEditorRestart")
-if reset_deps:
+reset_system_deps = dependency_options.get_editor_property("ClearSystemDependenciesOnEditorRestart")
+
+if reset_deps or reset_system_deps:
     print(f"Clearing python dependendencies")
-    dep_manager.clear_all_dependencies(env_dir)
+    dep_manager.clear_all_dependencies(env_dir, reset_system_deps)
 
     # Flag dependencies as cleared so we don't keep clearing them every restart
     dep_manager.finished_clearing_dependencies()

StableDiffusionTools/Content/Python/requirements.txt

@@ -1,46 +1,79 @@
-accelerate
-certifi
-charset_normalizer
-colorama
-compel
-cv2
-diffusers
-dotenv
-filelock
-ftfy
-functorch
-google
-grpc
-grpc_tools
-huggingface_hub
-idna
-importlib_metadata
-jinja2
-markupsafe
-mpmath
-networkx
-numpy
-nvfuser
-packaging
-PIL
-psutil
-pyparsing
-regex
-requests
-safetensors
-scipy
-scipy.libs
-setuptools
-share
-stability_sdk
-sympy
-tokenizers
-torch
-torchgen
-torchvision
-tqdm
-transformers
-urllib3
-wcwidth
-yaml
-zipp
+absl-py==1.4.0
+accelerate==0.18.0
+addict==2.4.0
+basicsr==1.4.2
+cachetools==5.3.0
+certifi==2022.12.7
+charset-normalizer==3.1.0
+colorama==0.4.6
+compel==1.1.3
+contourpy==1.0.7
+cycler==0.11.0
+diffusers==0.15.0
+facexlib==0.3.0
+filelock==3.11.0
+filterpy==1.4.5
+fonttools==4.39.3
+ftfy==6.1.1
+future==0.18.3
+gfpgan==1.3.8
+google-auth==2.17.3
+google-auth-oauthlib==1.0.0
+grpcio==1.53.0
+grpcio-tools==1.48.1
+huggingface-hub==0.13.4
+idna==3.4
+imageio==2.27.0
+importlib-metadata==6.4.1
+importlib-resources==5.12.0
+Jinja2==3.1.2
+kiwisolver==1.4.4
+lazy_loader==0.2
+llvmlite==0.39.1
+lmdb==1.4.1
+Markdown==3.4.3
+MarkupSafe==2.1.2
+matplotlib==3.7.1
+mpmath==1.3.0
+networkx==3.1
+numba==0.56.4
+numpy==1.24.2
+oauthlib==3.2.2
+opencv-python==4.7.0.72
+packaging==23.1
+Pillow==9.5.0
+protobuf==4.22.3
+psutil==5.9.4
+pyasn1==0.4.8
+pyasn1-modules==0.2.8
+pyparsing==3.0.9
+python-dateutil==2.8.2
+python-dotenv==1.0.0
+PyWavelets==1.4.1
+PyYAML==6.0
+realesrgan==0.3.0
+regex==2023.3.23
+requests==2.28.2
+requests-oauthlib==1.3.1
+rsa==4.9
+safetensors==0.3.0
+scikit-image==0.20.0
+scipy==1.10.1
+six==1.16.0
+stability-sdk==0.5.0
+sympy==1.11.1
+tb-nightly==2.13.0a20230414
+tensorboard-data-server==0.7.0
+tensorboard-plugin-wit==1.8.1
+tifffile==2023.4.12
+tokenizers==0.13.3
+torch==2.0.0+cu118
+torchvision==0.15.1+cu118
+tqdm==4.65.0
+transformers==4.28.1
+typing_extensions==4.5.0
+urllib3==1.26.15
+wcwidth==0.2.6
+Werkzeug==2.2.3
+yapf==0.32.0
+zipp==3.15.0

StableDiffusionTools/Content/UI/Widgets/BP_BridgeManifestRow.uasset

@@ -25,18 +25,20 @@ void UDependencyManager::RestartAndUpdateDependencies()
 	UStableDiffusionBlueprintLibrary::RestartEditor();
 }
 
-void UDependencyManager::ResetDependencies()
+void UDependencyManager::ResetDependencies(bool ClearSystemDeps)
 {
 	auto Settings = GetMutableDefault<UDependencySettings>();
 	Settings->ClearDependenciesOnEditorRestart = true;
+	Settings->ClearSystemDependenciesOnEditorRestart = ClearSystemDeps;
 	Settings->SaveConfig();
-	RestartAndUpdateDependencies();
+	//RestartAndUpdateDependencies();
 }
 
 void UDependencyManager::FinishedClearingDependencies()
 {
 	auto Settings = GetMutableDefault<UDependencySettings>();
 	Settings->ClearDependenciesOnEditorRestart = false;
+	Settings->ClearSystemDependenciesOnEditorRestart = false;
 	Settings->SaveConfig();
 }
 
@@ -45,6 +47,7 @@ void UDependencyManager::FinishedUpdatingDependencies()
 	auto Settings = GetMutableDefault<UDependencySettings>();
 	Settings->AutoLoadBridgeScripts = true;
 	Settings->ClearDependenciesOnEditorRestart = false;
+	Settings->ClearSystemDependenciesOnEditorRestart = false;
 	Settings->AutoUpdateDependenciesOnStartup = false;
 	Settings->SaveConfig();
 	UStableDiffusionBlueprintLibrary::RestartEditor();

StableDiffusionTools/Content/UI/Widgets/BP_StableDiffusionDependencyInstallerWidget.uasset

@@ -97,7 +97,7 @@ class STABLEDIFFUSIONTOOLS_API UDependencyManager : public UObject
         void RestartAndUpdateDependencies();
 
     UFUNCTION(BlueprintCallable, Category = "StableDiffusion|Dependencies")
-        void ResetDependencies();
+        void ResetDependencies(bool ClearSystemDeps = false);
 
     UFUNCTION(BlueprintCallable, Category = "StableDiffusion|Dependencies")
         void FinishedClearingDependencies();

StableDiffusionTools/Source/StableDiffusionTools/Private/DependencyManager.cpp

@@ -31,4 +31,10 @@ class STABLEDIFFUSIONTOOLS_API UDependencySettings : public UObject
 	**/
 	UPROPERTY(BlueprintReadWrite, config, EditAnywhere, Category = "Options")
 		bool ClearDependenciesOnEditorRestart = false;
+
+	/**
+	* Remove all downloaded python dependencies from Unreal's site-packages folder when the editor next starts.
+	**/
+	UPROPERTY(BlueprintReadWrite, config, EditAnywhere, Category = "Options")
+		bool ClearSystemDependenciesOnEditorRestart = false;
 };