cpu speed up

app.py

@@ -12,18 +12,19 @@ from diffusers import UniPCMultistepScheduler
 from diffusers import AutoencoderKL
 from diffusers import StableDiffusionPipeline
 from diffusers.loaders import LoraLoaderMixin
+import intel_extension_for_pytorch as ipex
 
 import os
 from os.path import join as opj
 
-token = os.getenv("ACCESS_TOKEN")
-os.system(f"python -m pip install git+https://{token}@github.com/logn-2024/StableGarment.git")
+# token = os.getenv("ACCESS_TOKEN")
+# os.system(f"python -m pip install git+https://{token}@github.com/logn-2024/StableGarment.git")
 
 from stablegarment.models import GarmentEncoderModel,ControlNetModel
 from stablegarment.piplines import StableGarmentPipeline,StableGarmentControlNetPipeline
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float32 if device=="cpu" else torch.float16
+device = "cpu" #"cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.bfloat16 if device=="cpu" else torch.float16
 height = 512
 width = 384
 
@@ -39,6 +40,27 @@ pipeline_t2i = StableGarmentPipeline.from_pretrained(base_model_path, vae=vae, t
 # pipeline = StableDiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V4.0_noVAE", vae=vae, torch_dtype=torch_dtype).to(device=device)
 pipeline_t2i.scheduler = scheduler
 
+
+# speed up for cpu
+# to channels last
+pipeline_t2i.unet = pipeline_t2i.unet.to(memory_format=torch.channels_last)
+pipeline_t2i.vae = pipeline_t2i.vae.to(memory_format=torch.channels_last)
+pipeline_t2i.text_encoder = pipeline_t2i.text_encoder.to(memory_format=torch.channels_last)
+# pipeline_t2i.safety_checker = pipeline_t2i.safety_checker.to(memory_format=torch.channels_last)
+
+# Create random input to enable JIT compilation
+sample = torch.randn(2,4,64,64).type(torch_dtype)
+timestep = torch.rand(1)*999
+encoder_hidden_status = torch.randn(2,77,768).type(torch_dtype)
+input_example = (sample, timestep, encoder_hidden_status)
+
+# optimize with IPEX
+pipeline_t2i.unet = ipex.optimize(pipeline_t2i.unet.eval(), dtype=torch.bfloat16, inplace=True, sample_input=input_example)
+pipeline_t2i.vae = ipex.optimize(pipeline_t2i.vae.eval(), dtype=torch.bfloat16, inplace=True)
+pipeline_t2i.text_encoder = ipex.optimize(pipeline_t2i.text_encoder.eval(), dtype=torch.bfloat16, inplace=True)
+# pipeline_t2i.safety_checker = ipex.optimize(pipeline_t2i.safety_checker.eval(), dtype=torch.bfloat16, inplace=True)
+
+
 pipeline_tryon = None
 '''
 # not ready

requirements.txt

@@ -6,4 +6,5 @@ spaces
 Pillow
 numpy
 accelerate
+intel_extension_for_pytorch
 # git+https://github.com/logn-2024/StableGarment.git