Initial commit

requirements.txt

@@ -1,4 +1,3 @@
-https://github.com/TheLastBen/fast-stable-diffusion/raw/main/precompiled/T4/xformers-0.0.13.dev0-py3-none-any.whl
 accelerate==0.12.0
 diffusers==0.9.0
 tokenizers==0.13.2

src/utils/gradio_utils.py

@@ -19,6 +19,8 @@ from transformers import CLIPTextModel, CLIPTokenizer
 from torch import autocast
 from src.diffusers_ import StableDiffusionPipeline
 
+from diffusers.utils.import_utils import is_xformers_available
+from packaging import version
 
 
 def launch_source():
@@ -72,9 +74,23 @@ def launch_optimize(img_in_real, prompt, n_hiper):
     CLIP_text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name, subfolder="text_encoder")#, use_auth_token=True)
     vae = AutoencoderKL.from_pretrained(pretrained_model_name, subfolder="vae")#, use_auth_token=True)
     unet = UNet2DConditionModel.from_pretrained(pretrained_model_name, subfolder="unet")#, use_auth_token=True)
+    unet.enable_xformers_memory_efficient_attention()
     noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
     
 
+    if is_xformers_available():
+        import xformers
+
+        xformers_version = version.parse(xformers.__version__)
+        if xformers_version == version.parse("0.0.16"):
+            print(
+                "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
+            )
+        unet.enable_xformers_memory_efficient_attention()
+    else:
+        raise ValueError("xformers is not available. Make sure it is installed correctly")
+        
+
     # Encode the input image.
     vae.to(accelerator.device, dtype=weight_dtype)
     input_image = img_in_real.convert("RGB")