app.py

import gradio as gr
from gradio_client import Client
import os
import logging
import requests
from urllib.parse import urlparse, parse_qs

webhook_server = os.getenv('webhook_server')

# Initialize the client for image generation
client_image = Client("mukaist/DALLE-4K")

# Define resolutions and default style
resolutions = {
    "896x1152": (896, 1152),
    "1024x1024": (1024, 1024),
    "1216x832": (1216, 832)
}
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
    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,
            negative_prompt="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
            use_negative_prompt=True,
            style=style,
            seed=0,
            width=width,
            height=height,
            guidance_scale=5,
            randomize_seed=True,
            api_name="/run"
        )
        logger.info("Image generation successful.")
        return result
    except Exception as e:
        logger.error(f"Error generating image: {e}")
        return None

def request_otp(user_chat_id):
    try:
        # Make the request to the /request_otp endpoint
        response = requests.post(f"{webhook_server}/request_otp", json={"user_chat_id": user_chat_id})
        response.raise_for_status()  # Ensure we raise an exception for HTTP errors
        
        # Log the response details
        logger.debug(f"Response status code: {response.status_code}")
        logger.debug(f"Response content: {response.text}")
        
        # Parse and return the JSON response
        response_data = response.json()
        
        # Check if the response contains an error
        if 'error' in response_data:
            logger.error(f"Error response from /request_otp: {response_data['error']}")
            return {"success": False, "error": response_data.get("error", "Unknown error")}
        
        # Return success response
        return {"success": response_data.get("success", False)}
    
    except requests.RequestException as e:
        # Log request exceptions and return an error response
        logger.error(f"Request exception: {e}")
        return {"success": False, "error": "Request exception occurred"}
    
    except ValueError as e:
        # Handle JSON decode errors
        logger.error(f"JSON decode error: {e}")
        return {"success": False, "error": "Invalid JSON response"}

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
    if not verify_otp(user_chat_id, otp):
        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):
    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()
        logger.info(f"User with ID {user_chat_id} had their points updated by {points}. Webhook response: {response.status_code}")
    except requests.RequestException as e:
        logger.error(f"Error updating user points via webhook: {e}")

def get_user_points(user_chat_id):
    try:
        response = requests.get(f"{webhook_server}/get_points", params={'user_chat_id': user_chat_id})
        response.raise_for_status()
        data = response.json()
        if 'error' in data:
            return "User not found"
        return data.get('points', 0)
    except requests.RequestException as e:
        logger.error(f"Error retrieving user points: {e}")
        return "Error retrieving points"

def notify_webhook(user_chat_id):
    webhook_url = f"{webhook_server}/notify_update"
    payload = {'user_chat_id': user_chat_id}
    try:
        response = requests.post(webhook_url, json=payload)
        response.raise_for_status()
        logger.info(f"Webhook notification sent successfully. Response: {response.status_code}")
    except requests.RequestException as e:
        logger.error(f"Error sending webhook notification: {e}")

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):
    if not user_chat_id.strip():
        return None, "User Chat ID cannot be blank. Please provide a valid User Chat ID."

    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."
        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)
    if isinstance(points, str):
        if points == "User not found":
            return None, "Please register with @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):
    otp_response = request_otp(user_chat_id)
    if otp_response.get("success"):
        return "OTP has been sent to your registered contact."
    return "Error requesting OTP. Please try again later."

def create_gradio_interface():
    with gr.Blocks() as interface:
        gr.HTML("""
        <h1>AI Image Generation DALL·E 4K</h1>
        <h3>Welcome! Please enter your user chat ID to continue.</h3>
        """)

        user_chat_id_input = gr.Textbox(label="Your UID", placeholder="Your UID")
        user_otp_input = gr.Textbox(label="OTP", type="password")
        
        with gr.Row():
            request_otp_button = gr.Button("Request OTP")
            user_otp_input
        
        points_output = gr.Textbox(label="Your Balances", placeholder="Your Points", interactive=False)
        
        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")
        
        result_output = gr.Image(label="Generated Image", type="pil")
        message_output = gr.Textbox(label="Result", placeholder="Results will be shown here", interactive=False)

        with gr.Row():
            user_chat_id_input
            points_output
            request_otp_button
            get_points_button
            
        with gr.Row():
            prompt_input
            resolution_dropdown
            generate_button

        generate_button.click(
            fn=handle_generate_image,
            inputs=[prompt_input, resolution_dropdown, user_chat_id_input],
            outputs=[result_output, message_output]
        )

        get_points_button.click(
            fn=retrieve_user_points,
            inputs=[user_chat_id_input, user_otp_input],
            outputs=[points_output]
        )
        
        request_otp_button.click(
            fn=request_otp_and_update_ui,
            inputs=[user_chat_id_input],
            outputs=[message_output]
        )
        
        gr.HTML("""
        <style>
            footer.svelte-1rjryqp {
                display: none !important;
            }
        </style>
        """)

    return interface

def launch_gradio():
    url = 'https://measmonysuon-imagen.hf.space/?user_chat_id=00000000'
    user_chat_id = extract_user_chat_id_from_url(url)
    
    interface = create_gradio_interface()
    interface.launch()

if __name__ == '__main__':
    launch_gradio()