app.py

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

# Base API URL for Hugging Face inference 
API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
API_TOKEN = os.getenv("HF_READ_TOKEN")
headers = {"Authorization": f"Bearer {API_TOKEN}"}
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):
    # 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"Parameters - Steps: {steps}, CFG Scale: {cfg_scale}, Seed: {seed}, Strength: {strength}, Width: {width}, Height: {height}")

    # 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

    # 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")])
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    print(f"Selected API token: {API_TOKEN}")  # Debug log

    # 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 == '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 == '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}"
        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

    # Make a request to the API to generate the image
    try:
        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:
        # Attempt to read the image from the response 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

css = """
footer {
    visibility: hidden;
}
"""

print("Initializing Gradio interface...")

# Define the Gradio interface
with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
    gr.Markdown("# AI Image Generator")
    
    with gr.Row():
        with gr.Column(scale=2):
            # Main prompt input
            text_prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe what you want to create...",
                lines=3
            )
            
            # Negative prompt
            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",
                lines=2
            )

            # Custom LoRA input
            custom_lora = gr.Textbox(
                label="Custom LoRA Path (Optional)",
                placeholder="e.g., multimodalart/vintage-ads-flux",
                lines=1
            )

        with gr.Column(scale=1):
            # Image dimensions
            with gr.Group():
                gr.Markdown("### Image Settings")
                width = gr.Slider(label="Width", value=1024, minimum=512, maximum=1216, step=64)
                height = gr.Slider(label="Height", value=1024, minimum=512, maximum=1216, step=64)
                
            # Generation parameters
            with gr.Group():
                gr.Markdown("### Generation Parameters")
                steps = gr.Slider(label="Steps", value=35, minimum=1, maximum=100, step=1)
                cfg = gr.Slider(label="CFG Scale", value=7, minimum=1, maximum=20, step=0.5)
                strength = gr.Slider(label="Strength", value=0.7, minimum=0, maximum=1, step=0.1)
                seed = gr.Slider(label="Seed (-1 for random)", value=-1, minimum=-1, maximum=1000000000, step=1)

# Model selection
            with gr.Group():
                gr.Markdown("### Model Selection")
                model_search = gr.Textbox(
                    label="Search Models",
                    placeholder="Type to filter models...",
                    lines=1
                )
                
                # Updated model list (reordered by popularity/recency)
                models_list = [
                    "Stable Diffusion 3.5 Large",
                    "Stable Diffusion 3.5 Large Turbo",
                    "Stable Diffusion XL",
                    "FLUX.1 [Schnell]",
                    "FLUX.1 [Dev]",
                    "Midjourney",
                    "DreamPhotoGASM",
                    "Disney",
                    "Leonardo AI Style Illustration",
                    "AbsoluteReality 1.8.1",
                    "Analog Redmond",
                    "Stable Diffusion 3 Medium",
                    "Flux Super Realism LoRA",
                    "Flux Realism LoRA",
                    "Selfie Photography",
                    "Character Design",
                    "Pixel Art XL",
                    "3D Sketchfab",
                    "Anime Collection", # Group of anime-related models
                    "Flux Animex V2",
                    "Flux Animeo V1",
                    "Flux AestheticAnime",
                    "90s Anime Art",
                    "Softserve Anime",
                    "Artistic Styles", # Group of artistic style models
                    "Brain Melt Acid Art",
                    "Retro Comic Flux",
                    "Purple Dreamy",
                    "SoftPasty Flux",
                    "Specialized", # Group of specialized models
                    "Flux Logo Design",
                    "Product Design",
                    "Propaganda Poster",
                    "Movie Board",
                    "Collage Flux",
                    # Additional models...
                ]

                model = gr.Radio(
                    label="Select Model",
                    choices=models_list,
                    value="Stable Diffusion 3.5 Large",
                    interactive=True
                )

                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)

                model_search.change(filter_models, inputs=model_search, outputs=model)

    # Generate button and output
    with gr.Row():
        generate_btn = gr.Button("Generate Image", variant="primary", size="lg")
        
    with gr.Row():
        image_output = gr.Image(
            type="pil",
            label="Generated Image",
            show_label=True
        )

    # Set up the generation event
    generate_btn.click(
        fn=query,
        inputs=[
            text_prompt,
            model,
            custom_lora,
            negative_prompt,
            steps,
            cfg,
            method,
            seed,
            strength,
            width,
            height
        ],
        outputs=image_output
    )


print("Launching Gradio interface...")
dalle.launch(show_api=False, share=False)