Update app.py

app.py

@@ -16,7 +16,7 @@ API_TOKEN = os.getenv('ZeroGPU')
 # Set your token
 login(token=API_TOKEN, add_to_git_credential=True)
 # Initialize API client
-api = HfApi(token=HF_TOKEN)
+api = HfApi(token=API_TOKEN)
 
 # Initialize the client for image generation "prithivMLmods/IMAGINEO-4K" "mukaist/DALLE-4K"
 client_image = Client("measmonysuon/DALLE-4K")
@@ -39,7 +39,6 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 def generate_image(prompt, resolution_key, style=DEFAULT_STYLE):
-    api.request_space_hardware(repo_id='measmonysuon/DALLE-4K', hardware='a10g-large')
     resolution = resolutions.get(resolution_key, (1024, 1024))
     width, height = resolution
     full_prompt = f"{prompt}, Canon EOS R5, 4K, Photo-Realistic, appearing photorealistic with super fine details, high resolution, natural look, hyper realistic photography, cinematic lighting, --ar 64:37, --v 6.0, --style raw, --stylize 750"
@@ -58,11 +57,9 @@ def generate_image(prompt, resolution_key, style=DEFAULT_STYLE):
             api_name="/run"
         )
         logger.info("Image generation successful.")
-        api.request_space_hardware(repo_id='measmonysuon/DALLE-4K', hardware='zero-a10g')
         return result
     except Exception as e:
         logger.error(f"Error generating image: {e}")
-        api.request_space_hardware(repo_id='measmonysuon/DALLE-4K', hardware='zero-a10g')
         return None
 
 def request_otp(user_chat_id):
@@ -144,9 +141,13 @@ def notify_webhook(user_chat_id):
     except requests.RequestException as e:
         logger.error(f"Error sending webhook notification: {e}")
 
-def notify_webhook_with_file(user_chat_id, full_file_url):
+def notify_webhook_with_file(user_chat_id, full_file_url, gpu_info):
     webhook_url = f"{webhook_server}/notify_image"
-    payload = {'user_chat_id': user_chat_id, 'file_path': full_file_url}
+    payload = {
+        'user_chat_id': user_chat_id,
+        'file_path': full_file_url,
+        'gpu_info': gpu_info
+    }
     try:
         response = requests.post(webhook_url, json=payload)
         response.raise_for_status()
@@ -172,21 +173,40 @@ def gradio_interface(prompt, resolution_key, user_chat_id):
 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')
 
-    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
+    # 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."
             
-            if os.path.exists(file_path):
-                update_user_points(user_chat_id, -5)
-                notify_webhook_with_file(user_chat_id, full_file_url)  # Notify webhook with full file URL
-                return full_file_url, "Check your telegram, The image was generated successfully."
             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."
+        
+        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='zero-a10g')
     
 def handle_generate_image(prompt, resolution_key, user_chat_id):
     points = get_user_points(user_chat_id)