app.py

import gradio as gr
import requests
import io
import random
import os
import time
from PIL import Image
import json

# Project by Nymbo

# Base API URL for Hugging Face inference
API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
# Retrieve the API token from environment variables
API_TOKEN = os.getenv("HF_READ_TOKEN")
headers = {"Authorization": f"Bearer {API_TOKEN}"}
# Timeout for requests
timeout = 100

def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024, provider="hf-inference", custom_api_key=""):
    # Debug log to indicate function start
    print("Starting query function...")
    # Print the parameters for debugging purposes
    print(f"Prompt: {prompt}")
    print(f"Model: {model}")
    print(f"Custom LoRA: {custom_lora}")
    print(f"Provider: {provider}")
    print(f"Parameters - Steps: {steps}, CFG Scale: {cfg_scale}, Seed: {seed}, Strength: {strength}, Width: {width}, Height: {height}")
    print(f"Custom API Key provided: {bool(custom_api_key.strip())}")  # Log whether a custom key was provided without printing the key

    # Check if the prompt is empty or None
    if prompt == "" or prompt is None:
        print("Prompt is empty or None. Exiting query function.")  # Debug log
        return None

    # Generate a unique key for tracking the generation process
    key = random.randint(0, 999)
    print(f"Generated key: {key}")  # Debug log

    # Determine which token to use - custom API key if provided, otherwise one of the environment variables
    if custom_api_key.strip() != "":
        print("USING CUSTOM API KEY: BYOK token provided by user is being used for authentication")
        API_TOKEN = custom_api_key.strip()
    else:
        # If no custom key, check if provider is "hf-inference"
        if provider != "hf-inference":
            print(f"ERROR: Custom API key is required for {provider} provider")
            raise gr.Error(f"A custom API key is required when using the {provider} provider. Please enter your key in the BYOK field.")
        
        # For hf-inference, randomly select an API token from available options to distribute the load
        API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN"), os.getenv("HF_READ_TOKEN_2"), os.getenv("HF_READ_TOKEN_3"), os.getenv("HF_READ_TOKEN_4"), os.getenv("HF_READ_TOKEN_5")])
        print("USING DEFAULT API KEY: Random environment variable token is being used for authentication")
    
    headers = {"Authorization": f"Bearer {API_TOKEN}"}

    # Enhance the prompt with additional details for better quality
    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
    print(f'Generation {key}: {prompt}')  # Debug log

    # Set the API URL based on the selected model or custom LoRA
    if custom_lora.strip() != "":
        API_URL = f"https://api-inference.huggingface.co/models/{custom_lora.strip()}"
    else:
        if model == 'Stable Diffusion XL':
            API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
        if model == 'FLUX.1 [Dev]':
            API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
        if model == 'FLUX.1 [Schnell]':
            API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
        if model == 'HiDream-I1-Full':
            API_URL = "https://api-inference.huggingface.co/models/HiDream-ai/HiDream-I1-Full"
        if model == 'HiDream-I1-Dev':
            API_URL = "https://api-inference.huggingface.co/models/HiDream-ai/HiDream-I1-Dev"
        if model == 'HiDream-I1-Fast':
            API_URL = "https://api-inference.huggingface.co/models/HiDream-ai/HiDream-I1-Fast"
        if model == 'Animagine 4.0':
            API_URL = "https://api-inference.huggingface.co/models/cagliostrolab/animagine-xl-4.0"
            prompt = f"masterpiece, high score, great score, absurdres, {prompt}"
        if model == 'Flux Icon Kit':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Icon-Kit-LoRA"
            prompt = f"Icon Kit, {prompt}"
        if model == 'Pixel Background':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Pixel-Background-LoRA"
            prompt = f"Pixel Background, {prompt}"
        if model == 'Meme XD':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Flux-Meme-Xd-LoRA"
            prompt = f"meme, {prompt}"
        if model == 'Chill Guy':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Flux-Chill-Guy-Zone"
            prompt = f"chill guy, {prompt}"
        if model == 'Pepe':
            API_URL = "https://api-inference.huggingface.co/models/openfree/pepe"
            prompt = f"pepe, {prompt}"
        if model == 'NSFWmodel':
            API_URL = "https://api-inference.huggingface.co/models/lexa862/NSFWmodel"
            prompt = f"nude, {prompt}"
        if model == 'Claude Art':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Claude-Art"
            prompt = f"claude art, {prompt}"
        if model == 'Open Genmoji':
            API_URL = "https://api-inference.huggingface.co/models/EvanZhouDev/open-genmoji"
            prompt = f"emoji, {prompt}"
        if model == 'EBook Creative Cover':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/EBook-Creative-Cover-Flux-LoRA"
            prompt = f"EBook Cover, {prompt}"
        if model == 'Flux Logo Design 2':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Logo-Design-Flux-LoRA"
            prompt = f"Logo Design, {prompt}"
        if model == 'Isometric 3D':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Isometric-3D-LoRA"
            prompt = f"Isometric 3D, {prompt}"
        if model == 'Flux Condensation':
            API_URL = "https://api-inference.huggingface.co/models/fofr/flux-condensation"
            prompt = f"CONDENSATION, {prompt}"
        if model == 'Flux Handwriting':
            API_URL = "https://api-inference.huggingface.co/models/fofr/flux-handwriting"
            prompt = f"HWRIT handwriting, {prompt}"
        if model == 'Shou Xin':
            API_URL = "https://api-inference.huggingface.co/models/Datou1111/shou_xin"
            prompt = f"shou_xin, pencil sketch, {prompt}"
        if model == 'Sketch Smudge':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Sketch-Smudge-LoRA"
            prompt = f"Sketch Smudge, {prompt}"
        if model == '80s Cyberpunk':
            API_URL = "https://api-inference.huggingface.co/models/fofr/flux-80s-cyberpunk"
            prompt = f"80s cyberpunk, {prompt}"
        if model == 'Coloring Book Flux':
            API_URL = "https://api-inference.huggingface.co/models/renderartist/coloringbookflux"
            prompt = f"c0l0ringb00k, coloring book, coloring book page, {prompt}"
        if model == 'Flux Miniature LoRA':
            API_URL = "https://api-inference.huggingface.co/models/gokaygokay/Flux-Miniature-LoRA"
            prompt = f"MNTR, miniature drawing, {prompt}"
        if model == 'Sketch Paint':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Sketch-Paint"
            prompt = f"Sketch paint, {prompt}"
        if model == 'Flux UltraRealism 2.0':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Canopus-LoRA-Flux-UltraRealism-2.0"
            prompt = f"Ultra realistic, {prompt}"
        if model == 'Midjourney Mix':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Midjourney-Mix-LoRA"
            prompt = f"midjourney mix, {prompt}"
        if model == 'Midjourney Mix 2':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Midjourney-Mix2-LoRA"
            prompt = f"MJ v6, {prompt}"
        if model == 'Flux Logo Design':
            API_URL = "https://api-inference.huggingface.co/models/Shakker-Labs/FLUX.1-dev-LoRA-Logo-Design"
            prompt = f"wablogo, logo, Minimalist, {prompt}"
        if model == 'Flux Uncensored':
            API_URL = "https://api-inference.huggingface.co/models/enhanceaiteam/Flux-uncensored"
        if model == 'Flux Uncensored V2':
            API_URL = "https://api-inference.huggingface.co/models/enhanceaiteam/Flux-Uncensored-V2"
        if model == 'Flux Tarot Cards':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Ton618-Tarot-Cards-Flux-LoRA"
            prompt = f"Tarot card, {prompt}"
        if model == 'Pixel Art Sprites':
            API_URL = "https://api-inference.huggingface.co/models/sWizad/pokemon-trainer-sprites-pixelart-flux"
            prompt = f"a pixel image, {prompt}"
        if model == '3D Sketchfab':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Castor-3D-Sketchfab-Flux-LoRA"
            prompt = f"3D Sketchfab, {prompt}"
        if model == 'Retro Comic Flux':
            API_URL = "https://api-inference.huggingface.co/models/renderartist/retrocomicflux"
            prompt = f"c0m1c, comic book panel, {prompt}"
        if model == 'Caricature':
            API_URL = "https://api-inference.huggingface.co/models/TheAwakenOne/caricature"
            prompt = f"CCTUR3, {prompt}"
        if model == 'Huggieverse':
            API_URL = "https://api-inference.huggingface.co/models/Chunte/flux-lora-Huggieverse"
            prompt = f"HGGRE, {prompt}"
        if model == 'Propaganda Poster':
            API_URL = "https://api-inference.huggingface.co/models/AlekseyCalvin/Propaganda_Poster_Schnell_by_doctor_diffusion"
            prompt = f"propaganda poster, {prompt}"
        if model == 'Flux Game Assets V2':
            API_URL = "https://api-inference.huggingface.co/models/gokaygokay/Flux-Game-Assets-LoRA-v2"
            prompt = f"wbgmsst, white background, {prompt}"
        if model == 'SDXL HS Card Style':
            API_URL = "https://api-inference.huggingface.co/models/Norod78/sdxl-hearthstone-card-style-lora"
            prompt = f"Hearthstone Card, {prompt}"
        if model == 'SLDR FLUX NSFW v2 Studio':
            API_URL = "https://api-inference.huggingface.co/models/xey/sldr_flux_nsfw_v2-studio"
        if model == 'SoftPasty Flux':
            API_URL = "https://api-inference.huggingface.co/models/alvdansen/softpasty-flux-dev"
            prompt = f"araminta_illus illustration style, {prompt}"
        if model == 'Flux Stickers':
            API_URL = "https://api-inference.huggingface.co/models/diabolic6045/Flux_Sticker_Lora"
            prompt = f"5t1cker 5ty1e, {prompt}"
        if model == 'Flux Animex V2':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Animex-v2-LoRA"
            prompt = f"Animex, {prompt}"
        if model == 'Flux Animeo V1':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Animeo-v1-LoRA"
            prompt = f"Animeo, {prompt}"
        if model == 'Movie Board':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Flux.1-Dev-Movie-Boards-LoRA"
            prompt = f"movieboard, {prompt}"
        if model == 'Purple Dreamy':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Purple-Dreamy-Flux-LoRA"
            prompt = f"Purple Dreamy, {prompt}"
        if model == 'PS1 Style Flux':
            API_URL = "https://api-inference.huggingface.co/models/veryVANYA/ps1-style-flux"
            prompt = f"ps1 game screenshot, {prompt}"
        if model == 'Softserve Anime':
            API_URL = "https://api-inference.huggingface.co/models/alvdansen/softserve_anime"
            prompt = f"sftsrv style illustration, {prompt}"
        if model == 'Flux Tarot v1':
            API_URL = "https://api-inference.huggingface.co/models/multimodalart/flux-tarot-v1"
            prompt = f"in the style of TOK a trtcrd tarot style, {prompt}"
        if model == 'Half Illustration':
            API_URL = "https://api-inference.huggingface.co/models/davisbro/half_illustration"
            prompt = f"in the style of TOK, {prompt}"
        if model == 'OpenDalle v1.1':
            API_URL = "https://api-inference.huggingface.co/models/dataautogpt3/OpenDalleV1.1"
        if model == 'Flux Ghibsky Illustration':
            API_URL = "https://api-inference.huggingface.co/models/aleksa-codes/flux-ghibsky-illustration"
            prompt = f"GHIBSKY style, {prompt}"
        if model == 'Flux Koda':
            API_URL = "https://api-inference.huggingface.co/models/alvdansen/flux-koda"
            prompt = f"flmft style, {prompt}"
        if model == 'Soviet Diffusion XL':
            API_URL = "https://api-inference.huggingface.co/models/openskyml/soviet-diffusion-xl"
            prompt = f"soviet poster, {prompt}"
        if model == 'Flux Realism LoRA':
            API_URL = "https://api-inference.huggingface.co/models/XLabs-AI/flux-RealismLora"
        if model == 'Frosting Lane Flux':
            API_URL = "https://api-inference.huggingface.co/models/alvdansen/frosting_lane_flux"
            prompt = f"frstingln illustration, {prompt}"
        if model == 'Phantasma Anime':
            API_URL = "https://api-inference.huggingface.co/models/alvdansen/phantasma-anime"
        if model == 'Boreal':
            API_URL = "https://api-inference.huggingface.co/models/kudzueye/Boreal"
            prompt = f"photo, {prompt}"
        if model == 'How2Draw':
            API_URL = "https://api-inference.huggingface.co/models/glif/how2draw"
            prompt = f"How2Draw, {prompt}"
        if model == 'Flux AestheticAnime':
            API_URL = "https://api-inference.huggingface.co/models/dataautogpt3/FLUX-AestheticAnime"
        if model == 'Fashion Hut Modeling LoRA':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Fashion-Hut-Modeling-LoRA"
            prompt = f"Modeling of, {prompt}"
        if model == 'Flux SyntheticAnime':
            API_URL = "https://api-inference.huggingface.co/models/dataautogpt3/FLUX-SyntheticAnime"
            prompt = f"1980s anime screengrab, VHS quality, syntheticanime, {prompt}"
        if model == 'Flux Midjourney Anime':
            API_URL = "https://api-inference.huggingface.co/models/brushpenbob/flux-midjourney-anime"
            prompt = f"egmid, {prompt}"
        if model == 'Coloring Book Generator':
            API_URL = "https://api-inference.huggingface.co/models/robert123231/coloringbookgenerator"
        if model == 'Collage Flux':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Castor-Collage-Dim-Flux-LoRA"
            prompt = f"collage, {prompt}"
        if model == 'Flux Product Ad Backdrop':
            API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Flux-Product-Ad-Backdrop"
            prompt = f"Product Ad, {prompt}"
        if model == 'Product Design':
            API_URL = "https://api-inference.huggingface.co/models/multimodalart/product-design"
            prompt = f"product designed by prdsgn, {prompt}"
        if model == '90s Anime Art':
            API_URL = "https://api-inference.huggingface.co/models/glif/90s-anime-art"
        if model == 'Brain Melt Acid Art':
            API_URL = "https://api-inference.huggingface.co/models/glif/Brain-Melt-Acid-Art"
            prompt = f"maximalism, in an acid surrealism style, {prompt}"
        if model == 'Lustly Flux Uncensored v1':
            API_URL = "https://api-inference.huggingface.co/models/lustlyai/Flux_Lustly.ai_Uncensored_nsfw_v1"
        if model == 'NSFW Master Flux':
            API_URL = "https://api-inference.huggingface.co/models/Keltezaa/NSFW_MASTER_FLUX"
            prompt = f"NSFW, {prompt}"
        if model == 'Flux Outfit Generator':
            API_URL = "https://api-inference.huggingface.co/models/tryonlabs/FLUX.1-dev-LoRA-Outfit-Generator"
        if model == 'Midjourney':
            API_URL = "https://api-inference.huggingface.co/models/Jovie/Midjourney"
        if model == 'DreamPhotoGASM':
            API_URL = "https://api-inference.huggingface.co/models/Yntec/DreamPhotoGASM"
        if model == 'Flux Super Realism LoRA':
            API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Super-Realism-LoRA"
        if model == 'Stable Diffusion 2-1':
            API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2-1-base"
        if model == 'Stable Diffusion 3.5 Large':
            API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large"
        if model == 'Stable Diffusion 3.5 Large Turbo':
            API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large-turbo"
        if model == 'Stable Diffusion 3 Medium':
            API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3-medium-diffusers"
            prompt = f"A, {prompt}"
            
    # Adjust URL for different providers if needed
    base_url = "https://api-inference.huggingface.co"
    if provider == "fal-ai":
        base_url = "https://api.fal.ai/hf"
    elif provider == "nebius":
        base_url = "https://api.nebius.ai/hf"
    elif provider == "replicate":
        base_url = "https://api.replicate.com/v1/models"
        # Replicate has a different endpoint structure
        if "https://api-inference.huggingface.co/models/" in API_URL:
            model_id = API_URL.replace("https://api-inference.huggingface.co/models/", "")
            API_URL = f"{base_url}/{model_id}"
    elif provider == "together":
        base_url = "https://api.together.xyz/inference"
        if "https://api-inference.huggingface.co/models/" in API_URL:
            model_id = API_URL.replace("https://api-inference.huggingface.co/models/", "")
            API_URL = f"{base_url}/{model_id}"
    
    # Only update URL if we're not using a custom URL format (like for replicate)
    if provider != "replicate" and provider != "together" and "https://api-inference.huggingface.co/models/" in API_URL:
        model_id = API_URL.replace("https://api-inference.huggingface.co/models/", "")
        API_URL = f"{base_url}/models/{model_id}"
            
    print(f"API URL set to: {API_URL}")  # Debug log
        if model == 'Duchaiten Real3D NSFW XL':
            API_URL = "https://api-inference.huggingface.co/models/stablediffusionapi/duchaiten-real3d-nsfw-xl"
        if model == 'Pixel Art XL':
            API_URL = "https://api-inference.huggingface.co/models/nerijs/pixel-art-xl"
            prompt = f"pixel art, {prompt}"
        if model == 'Character Design':
            API_URL = "https://api-inference.huggingface.co/models/KappaNeuro/character-design"
            prompt = f"Character Design, {prompt}"
        if model == 'Sketched Out Manga':
            API_URL = "https://api-inference.huggingface.co/models/alvdansen/sketchedoutmanga"
            prompt = f"daiton, {prompt}"
        if model == 'Archfey Anime':
            API_URL = "https://api-inference.huggingface.co/models/alvdansen/archfey_anime"
        if model == 'Lofi Cuties':
            API_URL = "https://api-inference.huggingface.co/models/alvdansen/lofi-cuties"
        if model == 'YiffyMix':
            API_URL = "https://api-inference.huggingface.co/models/Yntec/YiffyMix"
        if model == 'Analog Madness Realistic v7':
            API_URL = "https://api-inference.huggingface.co/models/digiplay/AnalogMadness-realistic-model-v7"
        if model == 'Selfie Photography':
            API_URL = "https://api-inference.huggingface.co/models/artificialguybr/selfiephotographyredmond-selfie-photography-lora-for-sdxl"
            prompt = f"instagram model, discord profile picture, {prompt}"
        if model == 'Filmgrain':
            API_URL = "https://api-inference.huggingface.co/models/artificialguybr/filmgrain-redmond-filmgrain-lora-for-sdxl"
            prompt = f"Film Grain, FilmGrainAF, {prompt}"
        if model == 'Leonardo AI Style Illustration':
            API_URL = "https://api-inference.huggingface.co/models/goofyai/Leonardo_Ai_Style_Illustration"
            prompt = f"leonardo style, illustration, vector art, {prompt}"
        if model == 'Cyborg Style XL':
            API_URL = "https://api-inference.huggingface.co/models/goofyai/cyborg_style_xl"
            prompt = f"cyborg style, {prompt}"
        if model == 'Little Tinies':
            API_URL = "https://api-inference.huggingface.co/models/alvdansen/littletinies"
        if model == 'NSFW XL':
            API_URL = "https://api-inference.huggingface.co/models/Dremmar/nsfw-xl"
        if model == 'Analog Redmond':
            API_URL = "https://api-inference.huggingface.co/models/artificialguybr/analogredmond"
            prompt = f"timeless style, {prompt}"
        if model == 'Pixel Art Redmond':
            API_URL = "https://api-inference.huggingface.co/models/artificialguybr/PixelArtRedmond"
            prompt = f"Pixel Art, {prompt}"
        if model == 'Ascii Art':
            API_URL = "https://api-inference.huggingface.co/models/CiroN2022/ascii-art"
            prompt = f"ascii art, {prompt}"
        if model == 'Analog':
            API_URL = "https://api-inference.huggingface.co/models/Yntec/Analog"
        if model == 'Maple Syrup':
            API_URL = "https://api-inference.huggingface.co/models/Yntec/MapleSyrup"
        if model == 'Perfect Lewd Fantasy':
            API_URL = "https://api-inference.huggingface.co/models/digiplay/perfectLewdFantasy_v1.01"
        if model == 'AbsoluteReality 1.8.1':
            API_URL = "https://api-inference.huggingface.co/models/digiplay/AbsoluteReality_v1.8.1"
        if model == 'Disney':
            API_URL = "https://api-inference.huggingface.co/models/goofyai/disney_style_xl"
            prompt = f"Disney style, {prompt}"
        if model == 'Redmond SDXL':
            API_URL = "https://api-inference.huggingface.co/models/artificialguybr/LogoRedmond-LogoLoraForSDXL-V2"
        if model == 'epiCPhotoGasm':
            API_URL = "https://api-inference.huggingface.co/models/Yntec/epiCPhotoGasm"
    print(f"API URL set to: {API_URL}")  # Debug log

    # Define the payload for the request
    payload = {
        "inputs": prompt,
        "is_negative": is_negative,  # Whether to use a negative prompt
        "steps": steps,  # Number of sampling steps
        "cfg_scale": cfg_scale,  # Scale for controlling adherence to prompt
        "seed": seed if seed != -1 else random.randint(1, 1000000000),  # Random seed for reproducibility
        "strength": strength,  # How strongly the model should transform the image
        "parameters": {
            "width": width,  # Width of the generated image
            "height": height  # Height of the generated image
        }
    }
    print(f"Payload: {json.dumps(payload, indent=2)}")  # Debug log

    # Adjust payload for different providers if needed
    if provider == "replicate" or provider == "together":
        # These providers might have different API formats
        if provider == "replicate":
            payload = {
                "version": model_id,
                "input": {
                    "prompt": prompt,
                    "negative_prompt": is_negative,
                    "num_inference_steps": steps,
                    "guidance_scale": cfg_scale,
                    "seed": seed if seed != -1 else random.randint(1, 1000000000),
                    "strength": strength,
                    "width": width,
                    "height": height
                }
            }
        elif provider == "together":
            payload = {
                "model": model_id,
                "prompt": prompt,
                "negative_prompt": is_negative,
                "num_inference_steps": steps,
                "guidance_scale": cfg_scale,
                "seed": seed if seed != -1 else random.randint(1, 1000000000),
                "strength": strength,
                "width": width,
                "height": height
            }

    # Make a request to the API to generate the image
    try:
        # Log which provider we're using
        print(f"Sending request to {provider} provider")
        
        # Add provider to headers for HF inference if needed
        if provider != "hf-inference" and "api-inference.huggingface.co" in API_URL:
            headers["X-Provider"] = provider
            
        response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
        print(f"Response status code: {response.status_code}")  # Debug log
    except requests.exceptions.RequestException as e:
        # Log any request exceptions and raise an error for the user
        print(f"Request failed: {e}")  # Debug log
        raise gr.Error(f"Request failed: {e}")

    # Check if the response status is not successful
    if response.status_code != 200:
        print(f"Error: Failed to retrieve image. Response status: {response.status_code}")  # Debug log
        print(f"Response content: {response.text}")  # Debug log
        if response.status_code == 400:
            raise gr.Error(f"{response.status_code}: Bad Request - There might be an issue with the input parameters.")
        elif response.status_code == 401:
            raise gr.Error(f"{response.status_code}: Unauthorized - Please check your API token.")
        elif response.status_code == 403:
            raise gr.Error(f"{response.status_code}: Forbidden - You do not have permission to access this model.")
        elif response.status_code == 404:
            raise gr.Error(f"{response.status_code}: Not Found - The requested model could not be found.")
        elif response.status_code == 503:
            raise gr.Error(f"{response.status_code}: The model is being loaded. Please try again later.")
        else:
            raise gr.Error(f"{response.status_code}: An unexpected error occurred.")

    try:
        # Handle different response formats based on provider
        if provider == "replicate":
            # Replicate might return a URL to the image rather than the image itself
            result = response.json()
            if "output" in result and isinstance(result["output"], list) and len(result["output"]) > 0:
                image_url = result["output"][0]
                # Fetch the image from the URL
                image_response = requests.get(image_url, timeout=timeout)
                image_bytes = image_response.content
            else:
                raise Exception(f"Unexpected Replicate response format: {result}")
        else:
            # Standard response with image content
            image_bytes = response.content
            
        image = Image.open(io.BytesIO(image_bytes))
        print(f'Generation {key} completed! ({prompt})')  # Debug log
        return image
    except Exception as e:
        # Handle any errors that occur when opening the image
        print(f"Error while trying to open image: {e}")  # Debug log
        return None

# Custom CSS to hide the footer in the interface
css = """
* {}
footer {visibility: hidden !important;}
"""

print("Initializing Gradio interface...")  # Debug log

# Define the Gradio interface
with gr.Blocks(theme='Nymbo/Nymbo_Theme_5') as dalle:
    # Tab for basic settings
    with gr.Tab("Basic Settings"):
        with gr.Row():
            with gr.Column(elem_id="prompt-container"):
                with gr.Row():
                    # Textbox for user to input the prompt
                    text_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here", lines=3, elem_id="prompt-text-input")
                with gr.Row():
                    # Textbox for custom LoRA input
                    custom_lora = gr.Textbox(label="Custom LoRA", info="LoRA Hugging Face path (optional)", placeholder="multimodalart/vintage-ads-flux")
                with gr.Row():
                    # Accordion for selecting the model
                    with gr.Accordion("Featured Models", open=False):
                        # Textbox for searching models
                        model_search = gr.Textbox(label="Filter Models", placeholder="Search for a featured model...", lines=1, elem_id="model-search-input")
                        models_list = (
                            "3D Sketchfab",
                            "80s Cyberpunk",
                            "90s Anime Art",
                            "AbsoluteReality 1.8.1",
                            "Analog",
                            "Analog Madness Realistic v7",
                            "Analog Redmond",
                            "Animagine 4.0",
                            "Archfey Anime",
                            "Ascii Art",
                            "Brain Melt Acid Art",
                            "Boreal",
                            "Caricature",
                            "Collage Flux",
                            "Coloring Book Flux",
                            "Character Design",
                            "Chill Guy",
                            "Claude Art",
                            "Coloring Book Generator",
                            "Cyborg Style XL",
                            "Disney",
                            "DreamPhotoGASM",
                            "Duchaiten Real3D NSFW XL",
                            "EBook Creative Cover",
                            "EpiCPhotoGasm",
                            "Fashion Hut Modeling LoRA",
                            "Filmgrain",
                            "FLUX.1 [Dev]",
                            "FLUX.1 [Schnell]",
                            "FLux Condensation",
                            "Flux Handwriting",
                            "Flux Realism LoRA",
                            "Flux Super Realism LoRA",
                            "Flux Uncensored",
                            "Flux Uncensored V2",
                            "Flux Game Assets V2",
                            "Flux Icon Kit",
                            "Flux Ghibsky Illustration",
                            "Flux Animex V2",
                            "Flux Animeo V1",
                            "Flux AestheticAnime",
                            "Flux SyntheticAnime",
                            "Flux Stickers",
                            "Flux Koda",
                            "Flux Tarot v1",
                            "Flux Tarot Cards",
                            "Flux UltraRealism 2.0",
                            "Flux Midjourney Anime",
                            "Flux Miniature LoRA",
                            "Flux Logo Design",
                            "Flux Logo Design 2",
                            "Flux Product Ad Backdrop",
                            "Flux Outfit Generator",
                            "Frosting Lane Flux",
                            "Half Illustration",
                            "HiDream-I1-Full",
                            "HiDream-I1-Dev",
                            "HiDream-I1-Fast",
                            "How2Draw",
                            "Huggieverse",
                            "Isometric 3D",
                            "Leonardo AI Style Illustration",
                            "Little Tinies",
                            "Lofi Cuties",
                            "Lustly Flux Uncensored v1",
                            "Maple Syrup",
                            "Meme XD",
                            "Midjourney",
                            "Midjourney Mix",
                            "Midjourney Mix 2",
                            "Movie Board",
                            "NSFWmodel",
                            "NSFW Master Flux",
                            "NSFW XL",
                            "OpenDalle v1.1",
                            "Open Genmoji",
                            "Pepe",
                            "Perfect Lewd Fantasy",
                            "Pixel Art Redmond",
                            "Pixel Art XL",
                            "Pixel Art Sprites",
                            "Pixel Background",
                            "Product Design",
                            "Propaganda Poster",
                            "Purple Dreamy",
                            "Phantasma Anime",
                            "PS1 Style Flux",
                            "Redmond SDXL",
                            "Retro Comic Flux",
                            "SDXL HS Card Style",
                            "Sketch Smudge",
                            "Shou Xin",
                            "Softserve Anime",
                            "SoftPasty Flux",
                            "Soviet Diffusion XL",
                            "Sketched Out Manga",
                            "Sketch Paint",
                            "SLDR FLUX NSFW v2 Studio",
                            "Selfie Photography",
                            "Stable Diffusion 2-1",
                            "Stable Diffusion XL",
                            "Stable Diffusion 3 Medium",
                            "Stable Diffusion 3.5 Large",
                            "Stable Diffusion 3.5 Large Turbo",
                            "YiffyMix",
                        )

                        # Radio buttons to select the desired model
                        model = gr.Radio(label="Select a model below", value="FLUX.1 [Schnell]", choices=models_list, interactive=True, elem_id="model-radio")

                        # Filtering models based on search input
                        def filter_models(search_term):
                            filtered_models = [m for m in models_list if search_term.lower() in m.lower()]
                            return gr.update(choices=filtered_models)

                        # Update model list when search box is used
                        model_search.change(filter_models, inputs=model_search, outputs=model)

    # Tab for advanced settings
    with gr.Tab("Advanced Settings"):
        with gr.Row():
            # Textbox for specifying elements to exclude from the image
            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="What should not be in the image", value="(deformed, distorted, disfigured), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, misspellings, typos", lines=3, elem_id="negative-prompt-text-input")
        
        with gr.Row():
            # Provider selection radio buttons
            providers_list = [
                "hf-inference",  # Default Hugging Face Inference
                "fal-ai",        # Fal.ai
                "nebius",        # Nebius
                "replicate",     # Replicate
                "together",      # Together AI
            ]
            
            provider_radio = gr.Radio(
                choices=providers_list,
                value="hf-inference",
                label="Inference Provider",
                info="Select which provider to use for image generation. Providers other than HF Inference require a custom API key.",
                elem_id="provider-radio"
            )
        
        with gr.Row():
            # BYOK (Bring Your Own Key) textbox
            byok_textbox = gr.Textbox(
                value="",
                label="BYOK (Bring Your Own Key)",
                info="Enter a custom API key here. Required for all providers except HF Inference.",
                placeholder="Enter your API token",
                type="password",  # Hide the API key for security
                elem_id="byok-input"
            )
            
        # Function to validate provider selection based on BYOK
        def validate_provider(api_key, provider):
            # If no custom API key is provided, only "hf-inference" can be used
            if not api_key.strip() and provider != "hf-inference":
                return gr.update(value="hf-inference")
            return gr.update(value=provider)
            
        # Connect the BYOK textbox to validate provider selection
        byok_textbox.change(
            fn=validate_provider,
            inputs=[byok_textbox, provider_radio],
            outputs=provider_radio
        )
        
        # Also validate provider when the radio changes to ensure consistency
        provider_radio.change(
            fn=validate_provider,
            inputs=[byok_textbox, provider_radio],
            outputs=provider_radio
        )
        with gr.Row():
            # Slider for selecting the image width
            width = gr.Slider(label="Width", value=1024, minimum=64, maximum=1216, step=32)
            # Slider for selecting the image height
            height = gr.Slider(label="Height", value=1024, minimum=64, maximum=1216, step=32)
        with gr.Row():
            # Slider for setting the number of sampling steps
            steps = gr.Slider(label="Sampling steps", value=35, minimum=1, maximum=100, step=1)
        with gr.Row():
            # Slider for adjusting the CFG scale (guidance scale)
            cfg = gr.Slider(label="CFG Scale", value=7, minimum=1, maximum=20, step=1)
        with gr.Row():
            # Slider for adjusting the transformation strength
            strength = gr.Slider(label="Strength", value=0.7, minimum=0, maximum=1, step=0.001)
        with gr.Row():
            # Slider for setting the seed for reproducibility
            seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1)
        with gr.Row():
            # Radio buttons for selecting the sampling method
            method = gr.Radio(label="Sampling method", value="DPM++ 2M Karras", choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"])

    # Tab for image editing options
    with gr.Tab("Image Editor"):
        # Function to simulate a delay for processing
        def sleep(im):
            print("Sleeping for 5 seconds...")  # Debug log
            time.sleep(5)
            return [im["background"], im["layers"][0], im["layers"][1], im["composite"]]

        # Function to return the composite image
        def predict(im):
            print("Predicting composite image...")  # Debug log
            return im["composite"]

        with gr.Blocks() as demo:
            with gr.Row():
                # Image editor component for user adjustments
                im = gr.ImageEditor(
                    type="numpy",
                    crop_size="1:1",  # Set crop size to a square aspect ratio
                )

    # Tab to provide information to the user
    with gr.Tab("Information"):
        with gr.Row():
            # Display a sample prompt for guidance
            gr.Textbox(label="Sample prompt", value="{prompt} | ultra detail, ultra elaboration, ultra quality, perfect.")

        # Accordion displaying featured models
        with gr.Accordion("Featured Models (WiP)", open=False):
            gr.HTML(
                """
            <p><a href="https://huggingface.co/models?inference=warm&pipeline_tag=text-to-image&sort=trending">See all available models</a></p>
            <table style="width:100%; text-align:center; margin:auto;">
                <tr>
                    <th>Model Name</th>
                    <th>Typography</th>
                    <th>Notes</th>
                </tr>
                <tr>
                    <td>FLUX.1 Dev</td>
                    <td>✅</td>
                    <td></td>
                </tr>
                <tr>
                    <td>FLUX.1 Schnell</td>
                    <td>✅</td>
                    <td></td>
                </tr>
                <tr>
                    <td>Stable Diffusion 3.5 Large</td>
                    <td>✅</td>
                    <td></td>
                </tr>
            </table>
            """
            )

        # Accordion providing an overview of advanced settings
        with gr.Accordion("Advanced Settings Overview", open=False):
            gr.Markdown(
            """
            ## Negative Prompt
            ###### This box is for telling the AI what you don't want in your images. Think of it as a way to avoid certain elements. For instance, if you don't want blurry images or extra limbs showing up, this is where you'd mention it.

            ## Inference Provider
            ###### Select which AI provider to use for generating images. Different providers may have different capabilities, pricing, and performance characteristics:
            ###### - HF Inference: Default Hugging Face inference API (uses application's default keys)
            ###### - Fal AI: Optimized for low latency and high throughput (requires your own API key)
            ###### - Nebius: Cloud provider with enterprise-grade infrastructure (requires your own API key)
            ###### - Replicate: Wide variety of models with flexible deployment (requires your own API key)
            ###### - Together: High-performance inference service (requires your own API key)

            ## BYOK (Bring Your Own Key)
            ###### This allows you to use your own API key instead of the default keys. Enter your key here for direct access to models using your account's permissions and rate limits. A custom key is required for all providers except HF Inference.

            ## Width & Height
            ###### These sliders allow you to specify the resolution of your image. Default value is 1024x1024, and maximum output is 1216x1216.

            ## Sampling Steps
            ###### Think of this like the number of brushstrokes in a painting. A higher number can give you a more detailed picture, but it also takes a bit longer. Generally, a middle-ground number like 35 is a good balance between quality and speed.

            ## CFG Scale
            ###### CFG stands for "Control Free Guidance." The scale adjusts how closely the AI follows your prompt. A lower number makes the AI more creative and free-flowing, while a higher number makes it stick closely to what you asked for. If you want the AI to take fewer artistic liberties, slide this towards a higher number. Just think "Control Freak Gauge".

            ## Sampling Method
            ###### This is the technique the AI uses to create your image. Each option is a different approach, like choosing between pencils, markers, or paint. You don't need to worry too much about this; the default setting is usually the best choice for most users.

            ## Strength
            ###### This setting is a bit like the 'intensity' knob. It determines how much the AI modifies the base image it starts with. If you're looking to make subtle changes, keep this low. For more drastic transformations, turn it up.

            ## Seed
            ###### You can think of the seed as a 'recipe' for creating an image. If you find a seed that gives you a result you love, you can use it again to create a similar image. If you leave it at -1, the AI will generate a new seed every time.

            ### Remember, these settings are all about giving you control over the image generation process. Feel free to experiment and see what each one does. And if you're ever in doubt, the default settings are a great place to start. Happy creating!
            """
            )

    # Row containing the 'Run' button to trigger the image generation
    with gr.Row():
        text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
    # Row for displaying the generated image output
    with gr.Row():
        image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")

    # Set up button click event to call the query function with provider and BYOK parameters
    text_button.click(query, inputs=[text_prompt, model, custom_lora, negative_prompt, steps, cfg, method, seed, strength, width, height, provider_radio, byok_textbox], outputs=image_output)

print("Launching Gradio interface...")  # Debug log
# Launch the Gradio interface without showing the API or sharing externally
dalle.launch(show_api=True, share=False)