import gradio as gr
from gradio_client import Client
import os
import logging
import sqlite3
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')
db_path = os.getenv('db_path') # Ensure you have this set up
# 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 in the database."""
with sqlite3.connect(db_path) as conn:
cursor = conn.cursor()
cursor.execute('''
INSERT INTO users (user_chat_id, points) VALUES (?, ?)
ON CONFLICT(user_chat_id) DO UPDATE SET points = points + ?
''', (user_chat_id, points, points))
conn.commit()
logger.info(f"User with ID {user_chat_id} had their points updated by {points}.")
def notify_webhook(user_chat_id):
"""Notify the webhook server about the points update."""
webhook_url = f"{webhook_server}/update_points"
payload = {'user_chat_id': user_chat_id, 'points': 10} # Adjust payload as necessary
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 get_user_points(user_chat_id):
"""Retrieve user points from the Flask server."""
webhook_url = f"{webhook_server}/get_points"
params = {'user_chat_id': user_chat_id}
try:
response = requests.get(webhook_url, params=params)
response.raise_for_status()
return response.json().get('points', "Error retrieving points")
except requests.RequestException as e:
logger.error(f"Error fetching user points: {e}")
return "Failed to retrieve user 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):
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, 10) # Award 10 points for each image generated
notify_webhook(user_chat_id) # Notify the webhook server
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 create_gradio_interface():
with gr.Blocks() as interface:
# Personalized HTML content
gr.HTML("""
Welcome! Please enter your user chat ID to continue.
""")
# Create input components
user_chat_id_input = gr.Textbox(label="User Chat ID", placeholder="Enter your user chat ID here...")
points_output = gr.Textbox(label="User Points", placeholder="User points will be displayed here", interactive=False)
# Create the button to get user points
get_points_button = gr.Button("Get 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", visible=False) # Start with the button hidden
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
def handle_get_points(user_chat_id):
points = get_user_points(user_chat_id)
if points >= 5:
return points, gr.Button.update(visible=True)
else:
return points, gr.Button.update(visible=False)
def handle_generate_image(prompt, resolution_key, user_chat_id):
points = get_user_points(user_chat_id)
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."
# 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=handle_get_points,
inputs=[user_chat_id_input],
outputs=[points_output, generate_button]
)
gr.HTML("""
""")
return interface
def launch_gradio():
interface = create_gradio_interface()
interface.launch()
if __name__ == '__main__':
# Launch Gradio
launch_gradio()