Update app.py

app.py

@@ -18,6 +18,10 @@ login(token=API_TOKEN, add_to_git_credential=True)
 # Initialize API client
 api = HfApi(token=API_TOKEN)
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Initialize the client for image generation "prithivMLmods/IMAGINEO-4K" "mukaist/DALLE-4K"
 client_image = Client("measmonysuon/DALLE-4K")
 
@@ -34,10 +38,6 @@ resolutions = {
 }
 DEFAULT_STYLE = "3840 x 2160"
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
 def generate_image(prompt, resolution_key, style=DEFAULT_STYLE):
     resolution = resolutions.get(resolution_key, (1024, 1024))
     width, height = resolution
@@ -61,7 +61,77 @@ def generate_image(prompt, resolution_key, style=DEFAULT_STYLE):
     except Exception as e:
         logger.error(f"Error generating image: {e}")
         return None
+        
+def gradio_interface(prompt, resolution_key, user_chat_id):
+    if not user_chat_id.strip():
+        return None, "User Chat ID cannot be blank. Please provide a valid User Chat ID or register with @supportBot to get UID."
+    
+    # Request large hardware for image generation
+    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"
+    }
+
+    try:
+        # Generate the image
+        result = generate_image(prompt, resolution_key)
+        
+        if result and result[0]:
+            file_path = result[0][0].get('image')
+            if file_path:
+                # Combine base URL with file path to create full URL
+                base_url = "https://measmonysuon-imagen.hf.space/file="
+                full_file_url = base_url + file_path
+                
+                if os.path.exists(file_path):
+                    # Update user points and notify webhook
+                    update_user_points(user_chat_id, -5)
+                    notify_webhook_with_file(user_chat_id, full_file_url, gpu_info)  # Notify webhook with full file URL
+                    
+                    return full_file_url, f"Check your telegram, The image was generated successfully. {gpu_info}"
+                
+                return None, "The image file is not available. Please try again later."
+            
+            return None, "The image file is not available. Please try again later."
+        
+        return None, "There was an error processing your photo. Please try again later."
+
+    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)
+    
+    if isinstance(points, str):
+        if points == "User not found":
+            return None, "Please register with https://t.me/myexpsupportBot"
+        return None, "Error retrieving points. Please try again later."
+
+    try:
+        points = int(points)
+    except ValueError:
+        return None, "Error processing points. Please try again later."
+
+    if points >= 5:
+        result = gradio_interface(prompt, resolution_key, user_chat_id)
+        if result[0]:
+            # Include the GPU info in the success message
+            gpu_info = result[1]  # Assuming gradio_interface returns GPU info in the second part of the tuple
+            return result[0], f"Check your telegram, The image was generated successfully. {gpu_info}"
+        
+        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 request_otp(user_chat_id):
     try:
         response = requests.post(f"{webhook_server}/request_otp", json={"user_chat_id": user_chat_id})
@@ -170,63 +240,6 @@ def gradio_interface(prompt, resolution_key, user_chat_id):
     return None, "There was an error processing your photo. Please try again later."
 '''
 
-def gradio_interface(prompt, resolution_key, user_chat_id):
-    if not user_chat_id.strip():
-        return None, "User Chat ID cannot be blank. Please provide a valid User Chat ID or register with @supportBot to get UID."
-    
-    # Request large hardware for image generation
-    api.request_space_hardware(repo_id='measmonysuon/DALLE-4K', hardware='a10g-large')
-
-    # Get GPU info
-    gpu_info = f"Is CUDA available: {torch.cuda.is_available()}"
-
-    try:
-        # Generate the image
-        result = generate_image(prompt, resolution_key)
-        
-        if result and result[0]:
-            file_path = result[0][0].get('image')
-            if file_path:
-                # Combine base URL with file path to create full URL
-                base_url = "https://measmonysuon-imagen.hf.space/file="
-                full_file_url = base_url + file_path
-                
-                if os.path.exists(file_path):
-                    # Update user points and notify webhook
-                    update_user_points(user_chat_id, -5)
-                    notify_webhook_with_file(user_chat_id, full_file_url, gpu_info)  # Notify webhook with full file URL
-                    
-                    return full_file_url, f"Check your telegram, The image was generated successfully. {gpu_info}"
-                
-                return None, "The image file is not available. Please try again later."
-            
-            return None, "The image file is not available. Please try again later."
-        
-        return None, "There was an error processing your photo. Please try again later."
-
-    finally:
-        # Always request to switch back hardware
-        api.request_space_hardware(repo_id='measmonysuon/DALLE-4K', hardware='cpu-basic')
-    
-def handle_generate_image(prompt, resolution_key, user_chat_id):
-    points = get_user_points(user_chat_id)
-    if isinstance(points, str):
-        if points == "User not found":
-            return None, "Please register with https://t.me/myexpsupportBot"
-        return None, "Error retrieving points. Please try again later."
-
-    try:
-        points = int(points)
-    except ValueError:
-        return None, "Error processing points. Please try again later."
-
-    if points >= 5:
-        result = gradio_interface(prompt, resolution_key, user_chat_id)
-        if result[0]:
-            return result[0], "The image was generated successfully."
-        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 request_otp_and_update_ui(user_chat_id):
     # Validate user_chat_id
     if not user_chat_id.strip():