Update app.py

app.py

@@ -28,7 +28,6 @@ client_image = Client("measmonysuon/DALLE-4K")
 # Get user information to confirm authentication
 user_info = whoami(token=API_TOKEN)
 print("User Info:", user_info)
-print(f"Is CUDA available: {torch.cuda.is_available()}")
 
 # Define resolutions and default style
 resolutions = {
@@ -70,11 +69,7 @@ def gradio_interface(prompt, resolution_key, user_chat_id):
     api.request_space_hardware(repo_id='measmonysuon/DALLE-4K', hardware='a10g-large')
 
     # Get GPU info
-    gpu_info = {
-        "is_cuda_available": torch.cuda.is_available(),
-        "device_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "N/A",
-        "total_memory": torch.cuda.get_device_properties(0).total_memory if torch.cuda.is_available() else "N/A"
-    }
+    gpu_info = get_gpu_info()
 
     try:
         # Generate the image
@@ -131,7 +126,24 @@ def handle_generate_image(prompt, resolution_key, user_chat_id):
         return None, "There was an error processing your photo. Please try again later."
     
     return None, "Insufficient points. Please get more points before generating an image."
-    
+
+def get_gpu_info():
+    if torch.cuda.is_available():
+        device_name = torch.cuda.get_device_name(0)
+        device_properties = torch.cuda.get_device_properties(0)
+        total_memory = f"{device_properties.total_memory // (1024 ** 2)} MB"
+        return {
+            'is_cuda_available': True,
+            'device_name': device_name,
+            'total_memory': total_memory
+        }
+    else:
+        return {
+            'is_cuda_available': False,
+            'device_name': 'N/A',
+            'total_memory': 'N/A'
+        }
+        
 def request_otp(user_chat_id):
     try:
         response = requests.post(f"{webhook_server}/request_otp", json={"user_chat_id": user_chat_id})