Update app.py

app.py

@@ -3,7 +3,7 @@ from gradio_client import Client
 import os
 import logging
 import requests
-from huggingface_hub import login, whoami
+from huggingface_hub import login, whoami, HfFolder, HfApi, SpaceHardware
 import torch
 
 webhook_server = os.getenv('webhook_server')
@@ -15,6 +15,8 @@ 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)
 
 # Initialize the client for image generation "prithivMLmods/IMAGINEO-4K" "mukaist/DALLE-4K"
 client_image = Client("measmonysuon/DALLE-4K")
@@ -37,6 +39,7 @@ 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"
@@ -55,9 +58,11 @@ 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):