Update app.py

app.py

@@ -22,6 +22,19 @@ from flux.util import (configs, embed_watermark, load_ae, load_clip, load_flow_m
 from huggingface_hub import login
 login(token=os.getenv('Token'))
 
+
+import torch
+
+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")
+
+
 @dataclass
 class SamplingOptions:
     source_prompt: str