app.py

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

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

# Retrieve secret token from environment variables
webhook_server = os.getenv('webhook_server')

# Define resolutions
resolutions = {
    "896x1152": (896, 1152),
    "1024x1024": (1024, 1024),
    "1216x832": (1216, 832)
}

# Define the default style
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 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
    }
    
    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):
    """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()
        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):
    """Notify the webhook server about the points update."""
    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 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 f"Your current points: {points}", 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."

def handle_generate_image(prompt, resolution_key, user_chat_id):
    points = get_user_points(user_chat_id)
    
    if points == "User not found":
        # If user is not found, update points_output with the registration message
        return None, "Please register with @myexpsupportBot"
    
    if points >= 5:
        result = gradio_interface(prompt, resolution_key, user_chat_id)
        if result[0]:  # If image generation is successful
            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."        
        

def create_gradio_interface():
    with gr.Blocks() as interface:
        # Personalized HTML content
        gr.HTML("""
        <h1> ImageGen DALL v3</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")
        
        # Create the button to get user points
        get_points_button = gr.Button("Show Your Points")
        
        # Create other components
        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)

        # Arrange components in rows
        with gr.Row():
            user_chat_id_input
            points_output
            get_points_button
            
        with gr.Row():
            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],
            outputs=[result_output, message_output]
        )

        get_points_button.click(
            fn=retrieve_user_points,
            inputs=[user_chat_id_input],
            outputs=[points_output]  # Update points output
        )

        gr.HTML("""
        <style>
            footer.svelte-1rjryqp {
                display: none !important;
            }
        </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
    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()

if __name__ == '__main__':
    launch_gradio()