update

app.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import devicetorch
 import gradio as gr
-import spaces
+#import spaces
 
 import argparse
 import inspect
@@ -1196,7 +1197,8 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             if use_md_prompt or save_attention_map:
                 self.recover_attention_control(ori_attn_processors=ori_attn_processors)  # recover attention controller
                 del self.controller
-            torch.cuda.empty_cache()
+            devicetorch.empty_cache(torch)
+            #torch.cuda.empty_cache()
         else:
             print("### Encoding Real Image ###")
             latents = self.vae.encode(image_lr)
@@ -1206,7 +1208,8 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
         anchor_std = latents.std()
         if self.lowvram:
             latents = latents.cpu()
-            torch.cuda.empty_cache()
+            #torch.cuda.empty_cache()
+            devicetorch.empty_cache(torch)
         if not output_type == "latent":
             # make sure the VAE is in float32 mode, as it overflows in float16
             needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
@@ -1227,7 +1230,8 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             # cast back to fp16 if needed
             if needs_upcasting:
                 self.vae.to(dtype=torch.float16)
-                torch.cuda.empty_cache()
+                #torch.cuda.empty_cache()
+                devicetorch.empty_cache(torch)
     
         image = self.image_processor.postprocess(image, output_type=output_type)
         if not os.path.exists(f'{result_path}'):
@@ -1250,7 +1254,8 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             if self.lowvram:
                 latents = latents.to(device)
                 self.unet.to(device)
-                torch.cuda.empty_cache()
+                #torch.cuda.empty_cache()
+                devicetorch.empty_cache(torch)
             
             current_height = self.unet.config.sample_size * self.vae_scale_factor * current_scale_num
             current_width = self.unet.config.sample_size * self.vae_scale_factor * current_scale_num
@@ -1549,7 +1554,8 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                 latents = (latents - latents.mean()) / latents.std() * anchor_std + anchor_mean
                 if self.lowvram:
                     latents = latents.cpu()
-                    torch.cuda.empty_cache()
+                    #torch.cuda.empty_cache()
+                    devicetorch.empty_cache(torch)
                 if not output_type == "latent":
                     # make sure the VAE is in float32 mode, as it overflows in float16
                     needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
@@ -1620,12 +1626,14 @@ if __name__ == "__main__":
 
     args = parser.parse_args()
     
-    pipe = AccDiffusionSDXLPipeline.from_pretrained(args.model_ckpt, torch_dtype=torch.float16).to("cuda")
+    #pipe = AccDiffusionSDXLPipeline.from_pretrained(args.model_ckpt, torch_dtype=torch.float16).to("cuda")
+    device = devicetorch.get(torch)
+    pipe = AccDiffusionSDXLPipeline.from_pretrained(args.model_ckpt, torch_dtype=torch.float16).to(device)
     
     
     # GRADIO MODE
 
-    @spaces.GPU()
+#    @spaces.GPU()
     def infer(prompt, resolution, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
         set_seed(seed)
         width,height = list(map(int, resolution.split(',')))
@@ -1634,7 +1642,8 @@ if __name__ == "__main__":
                                   "n_cross_replace": {"default_": 1.0, "confetti": 0.8},
                                   }
         seed = seed
-        generator = torch.Generator(device='cuda')
+        #generator = torch.Generator(device='cuda')
+        generator = torch.Generator(device=device)
         generator = generator.manual_seed(seed)
         
         print(f"Prompt: {prompt}")

requirements.txt

@@ -1,5 +1,5 @@
 diffusers~=0.21.4
-torch~=2.1.0
+#torch~=2.1.0
 scipy~=1.11.3
 omegaconf~=2.3.0
 accelerate~=0.23.0
@@ -10,4 +10,4 @@ matplotlib
 gradio
 gradio_imageslider
 opencv-python
-torchvision
+#torchvision