Update app.py

app.py

@@ -80,6 +80,8 @@ def gradio_interface(prompt, resolution_key, user_chat_id):
     # Add a short delay to allow the hardware switch to complete
     time.sleep(5)  # Adjust this value if needed
 
+    gpu_info = {}  # Initialize gpu_info to an empty dictionary
+    
     try:
         # Generate the image
         image_info, seed, gpu_info = generate_image(prompt, resolution_key)
@@ -111,6 +113,14 @@ def gradio_interface(prompt, resolution_key, user_chat_id):
         # Always request to switch back hardware
         api.request_space_hardware(repo_id='measmonysuon/DALLE-4K', hardware='zero-a10g')
         logger.info("Hardware reset to Zero-A10G")
+
+    finally:
+        # Log GPU info before resetting hardware
+        logger.info(f"GPU info before reset: {gpu_info}")
+        
+        # Always request to switch back hardware
+        api.request_space_hardware(repo_id='measmonysuon/DALLE-4K', hardware='zero-a10g')
+        logger.info("Hardware reset to Zero-A10G")
     
 def handle_generate_image(prompt, resolution_key, user_chat_id):
     points = get_user_points(user_chat_id)