Update app.py

app.py

@@ -25,16 +25,17 @@ login(token=os.getenv('Token'))
 
 import torch
 
-device = torch.cuda.current_device()
-print("!!!!!!!!!!!!device!!!!!!!!!!!!!!",device)
-total_memory = torch.cuda.get_device_properties(device).total_memory
-allocated_memory = torch.cuda.memory_allocated(device)
-reserved_memory = torch.cuda.memory_reserved(device)
+# device = torch.cuda.current_device()
+# print("!!!!!!!!!!!!device!!!!!!!!!!!!!!",device)
+# total_memory = torch.cuda.get_device_properties(device).total_memory
+# allocated_memory = torch.cuda.memory_allocated(device)
+# reserved_memory = torch.cuda.memory_reserved(device)
 
-print(f"Total memory: {total_memory / 1024**2:.2f} MB")
-print(f"Allocated memory: {allocated_memory / 1024**2:.2f} MB")
-print(f"Reserved memory: {reserved_memory / 1024**2:.2f} MB")
+# print(f"Total memory: {total_memory / 1024**2:.2f} MB")
+# print(f"Allocated memory: {allocated_memory / 1024**2:.2f} MB")
+# print(f"Reserved memory: {reserved_memory / 1024**2:.2f} MB")
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 name = 'flux-dev'
 ae = load_ae(name, device)
 t5 = load_t5(device, max_length=256 if name == "flux-schnell" else 512)
@@ -129,14 +130,14 @@ def edit(init_image, source_prompt, target_prompt, num_steps, inject_step, guida
     print("!!!!!!!!self.clip!!!!!!",next(clip.parameters()).device)
     print("!!!!!!!!self.model!!!!!!",next(model.parameters()).device)
 
-    device = torch.cuda.current_device()
-    total_memory = torch.cuda.get_device_properties(device).total_memory
-    allocated_memory = torch.cuda.memory_allocated(device)
-    reserved_memory = torch.cuda.memory_reserved(device)
+    # device = torch.cuda.current_device()
+    # total_memory = torch.cuda.get_device_properties(device).total_memory
+    # allocated_memory = torch.cuda.memory_allocated(device)
+    # reserved_memory = torch.cuda.memory_reserved(device)
         
-    print(f"Total memory: {total_memory / 1024**2:.2f} MB")
-    print(f"Allocated memory: {allocated_memory / 1024**2:.2f} MB")
-    print(f"Reserved memory: {reserved_memory / 1024**2:.2f} MB")
+    # print(f"Total memory: {total_memory / 1024**2:.2f} MB")
+    # print(f"Allocated memory: {allocated_memory / 1024**2:.2f} MB")
+    # print(f"Reserved memory: {reserved_memory / 1024**2:.2f} MB")
 
     with torch.no_grad():
         inp = prepare(t5, clip, init_image, prompt=opts.source_prompt)