Update app.py

app.py

@@ -9,16 +9,14 @@ from urllib.parse import urlparse, parse_qs
 client_image = Client("mukaist/DALLE-4K")
 
 # Retrieve secret token from environment variables
-webhook_server = os.getenv('webhook_server')
+webhook_server = os.getenv('WEBHOOK_SERVER')
 
-# Define resolutions
+# Define resolutions and default style
 resolutions = {
     "896x1152": (896, 1152),
     "1024x1024": (1024, 1024),
     "1216x832": (1216, 832)
 }
-
-# Define the default style
 DEFAULT_STYLE = "3840 x 2160"
 
 # Set up logging
@@ -29,7 +27,7 @@ def generate_image(prompt, resolution_key, style=DEFAULT_STYLE):
     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"
-    
+
     try:
         result = client_image.predict(
             prompt=full_prompt,
@@ -49,14 +47,27 @@ def generate_image(prompt, resolution_key, style=DEFAULT_STYLE):
         logger.error(f"Error generating image: {e}")
         return None
 
+def verify_otp(user_chat_id, otp):
+    response = requests.post(f"{webhook_server}/verify_otp", json={"user_chat_id": user_chat_id, "otp": otp})
+    return response.json().get("valid", False)
+
+def retrieve_user_points(user_chat_id, otp):
+    # Verify OTP via Flask server
+    otp_response = requests.post(f"{webhook_server}/verify_otp", json={"user_chat_id": user_chat_id, "otp": otp})
+    if otp_response.status_code != 200 or not otp_response.json().get('valid'):
+        return "Invalid or expired OTP. Please request a new OTP."
+
+    # Retrieve user points via Flask server
+    response = requests.get(f"{webhook_server}/get_points", params={"user_chat_id": user_chat_id})
+    if response.status_code == 200:
+        data = response.json()
+        points = data.get("points", 0)
+        return f"Your current balance is: {points}"
+    return "Error retrieving points. Please try again later."
+
 def update_user_points(user_chat_id, points):
-    """Updates the user's points using a webhook."""
-    webhook_url = f"{webhook_server}/update_user_points"  # Ensure this URL is correct
-    payload = {
-        'user_chat_id': user_chat_id,
-        'points': points
-    }
-    
+    webhook_url = f"{webhook_server}/update_user_points"
+    payload = {'user_chat_id': user_chat_id, 'points': points}
     try:
         response = requests.post(webhook_url, json=payload)
         response.raise_for_status()
@@ -65,7 +76,6 @@ def update_user_points(user_chat_id, points):
         logger.error(f"Error updating user points via webhook: {e}")
 
 def get_user_points(user_chat_id):
-    """Retrieve the user's points from the Flask server."""
     try:
         response = requests.get(f"{webhook_server}/get_points", params={'user_chat_id': user_chat_id})
         response.raise_for_status()
@@ -76,9 +86,8 @@ def get_user_points(user_chat_id):
     except requests.RequestException as e:
         logger.error(f"Error retrieving user points: {e}")
         return "Error retrieving points"
-    
+
 def notify_webhook(user_chat_id):
-    """Notify the webhook server about the points update."""
     webhook_url = f"{webhook_server}/notify_update"
     payload = {'user_chat_id': user_chat_id}
     try:
@@ -88,82 +97,56 @@ def notify_webhook(user_chat_id):
     except requests.RequestException as e:
         logger.error(f"Error sending webhook notification: {e}")
 
-def retrieve_user_points(user_chat_id):
-    """Retrieve user points and display them."""
-    points = get_user_points(user_chat_id)
-    if points == "Guest":
-        return "Please register with @myexpsupportBot", ""
-    return points
-    
 def extract_user_chat_id_from_url(url):
     parsed_url = urlparse(url)
     params = parse_qs(parsed_url.query)
     return params.get('user_chat_id', ["Guest"])[0]
 
 def gradio_interface(prompt, resolution_key, user_chat_id):
-    # Check if user_chat_id is empty
     if not user_chat_id.strip():
         return None, "User Chat ID cannot be blank. Please provide a valid User Chat ID."
 
-    # Proceed with image generation if user_chat_id is provided
     result = generate_image(prompt, resolution_key)
-    
     if result and result[0]:
         file_path = result[0][0].get('image')
         if file_path and os.path.exists(file_path):
             update_user_points(user_chat_id, -5)  # Deduct points for each image generated
             return file_path, "The image was generated successfully."
-        else:
-            return None, "The image file is not available. Please try again later."
-    else:
-        return None, "There was an error processing your photo. 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."
 
 def handle_generate_image(prompt, resolution_key, user_chat_id):
     points = get_user_points(user_chat_id)
-    
-    # Check if the points is a string indicating user not found
     if isinstance(points, str):
         if points == "User not found":
             return None, "Please register with @myexpsupportBot"
-        else:
-            return None, "Error retrieving points. Please try again later."
-    
-    # Ensure points is an integer
+        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]:  # If image generation is successful
+        if result[0]:
             return result[0], "The image was generated successfully."
-        else:
-            return None, "There was an error processing your photo. Please try again later."
-    else:
-        return None, "Insufficient points. Please get more points before generating an image."       
-        
+        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 create_gradio_interface():
     with gr.Blocks() as interface:
-        # Personalized HTML content
         gr.HTML("""
-        <h1> AI Image Generate DALL 4K</h1>
+        <h1>AI Image Generation DALL·E 4K</h1>
         <h3>Welcome! Please enter your user chat ID to continue.</h3>
         """)
 
-        # Create input components
-        user_chat_id_input = gr.Textbox(
-            label="Your UID",
-            placeholder="Your UID",
-            interactive=True  # Make this non-editable
-        )
-        points_output = gr.Textbox(label="Your Balances", placeholder="Your Points")
+        user_chat_id_input = gr.Textbox(label="Your UID", placeholder="Your UID")
+        user_otp_input = gr.Textbox(label="OTP", type="password")
+        points_output = gr.Textbox(label="Your Balances", placeholder="Your Points", interactive=False)
         
-        # Create the button to get user points
-        get_points_button = gr.Button("Show Your Points")
         
-        # Create other components
+        get_points_button = gr.Button("Show Your Points")
         prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
         resolution_dropdown = gr.Dropdown(choices=list(resolutions.keys()), label="Resolution", value="1024x1024")
         generate_button = gr.Button("Generate")
@@ -171,7 +154,6 @@ def create_gradio_interface():
         result_output = gr.Image(label="Generated Image", type="pil")
         message_output = gr.Textbox(label="Result", placeholder="Results will be shown here", interactive=False)
 
-        # Arrange components in rows
         with gr.Row():
             user_chat_id_input
             points_output
@@ -181,8 +163,7 @@ def create_gradio_interface():
             prompt_input
             resolution_dropdown
             generate_button
-        
-        # Set up interactions
+
         generate_button.click(
             fn=handle_generate_image,
             inputs=[prompt_input, resolution_dropdown, user_chat_id_input],
@@ -191,10 +172,10 @@ def create_gradio_interface():
 
         get_points_button.click(
             fn=retrieve_user_points,
-            inputs=[user_chat_id_input],
-            outputs=[points_output]  # Update points output
+            inputs=[user_chat_id_input, user_otp_input],
+            outputs=[points_output]
         )
-
+        
         gr.HTML("""
         <style>
             footer.svelte-1rjryqp {
@@ -202,15 +183,13 @@ def create_gradio_interface():
             }
         </style>
         """)
-    
+
     return interface
 
 def launch_gradio():
-    # Obtain the URL dynamically (from a configured environment variable or similar)
-    url = 'https://measmonysuon-imagen.hf.space/?user_chat_id=00000000'  # Replace with the actual URL or method to obtain it
+    url = 'https://measmonysuon-imagen.hf.space/?user_chat_id=00000000'
     user_chat_id = extract_user_chat_id_from_url(url)
     
-    # Pass the extracted user_chat_id to the Gradio interface
     interface = create_gradio_interface()
     interface.launch()