update

embodied_gen/models/sr_model.py

@@ -53,7 +53,7 @@ class ImageStableSR:
             torch_dtype=torch.float16,
         ).to(device)
         self.up_pipeline_x4.set_progress_bar_config(disable=True)
-        self.up_pipeline_x4.enable_model_cpu_offload()
+        # self.up_pipeline_x4.enable_model_cpu_offload()
 
     @spaces.GPU
     def __call__(

embodied_gen/models/text_model.py

@@ -92,7 +92,7 @@ def build_text2img_ip_pipeline(
     pipe.set_ip_adapter_scale([ref_scale])
 
     pipe = pipe.to(device)
-    pipe.enable_model_cpu_offload()
+    # pipe.enable_model_cpu_offload()
     # pipe.enable_xformers_memory_efficient_attention()
     # pipe.enable_vae_slicing()
 
@@ -123,8 +123,8 @@ def build_text2img_pipeline(
         force_zeros_for_empty_prompt=False,
     )
     pipe = pipe.to(device)
-    pipe.enable_model_cpu_offload()
-    pipe.enable_xformers_memory_efficient_attention()
+    # pipe.enable_model_cpu_offload()
+    # pipe.enable_xformers_memory_efficient_attention()
 
     return pipe
 

embodied_gen/models/texture_model.py

@@ -103,6 +103,6 @@ def build_texture_gen_pipe(
         pipe.set_ip_adapter_scale([ip_adapt_scale])
 
     pipe = pipe.to(device)
-    pipe.enable_model_cpu_offload()
+    # pipe.enable_model_cpu_offload()
 
     return pipe

requirements.txt

@@ -3,7 +3,7 @@ torchvision==0.19.0
 xformers==0.0.27.post2
 pytorch-lightning==2.4.0
 spconv-cu120==2.3.6
-triton==2.1.0
+triton
 dataclasses_json
 easydict
 opencv-python>4.5