app.py

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

# Load environment variables
load_dotenv()

# Get API token from environment variable 
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN environment variable is not set")

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):
    print("Starting query function...")
    
    if not prompt:
        raise gr.Error("Prompt cannot be empty")

    # Set headers with API token
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    
    # Generate a unique key for tracking
    key = random.randint(0, 999)
    
    # Enhance prompt
    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
    print(f'Generation {key}: {prompt}')

    # Set API URL based on model selection
    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"

    # Prepare payload
    payload = {
        "inputs": prompt,
        "is_negative": is_negative,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "seed": seed if seed != -1 else random.randint(1, 1000000000),
        "strength": strength,
        "parameters": {
            "width": width,
            "height": height
        }
    }

    try:
        response = requests.post(API_URL, headers=headers, json=payload, timeout=100)
        response.raise_for_status()
        
        image = Image.open(io.BytesIO(response.content))
        print(f'Generation {key} completed successfully')
        return image
        
    except requests.exceptions.RequestException as e:
        error_message = f"Request failed: {str(e)}"
        if response.status_code == 401:
            error_message = "Invalid API token. Please check your Hugging Face API token."
        elif response.status_code == 403:
            error_message = "Access denied. Please check your API token permissions."
        elif response.status_code == 503:
            error_message = "Model is currently loading. Please try again in a few moments."
        raise gr.Error(error_message)

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

print("Initializing 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):
            text_prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe what you want to create...",
                lines=3
            )
            
            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 = gr.Textbox(
                label="Custom LoRA Path (Optional)",
                placeholder="e.g., multimodalart/vintage-ads-flux",
                lines=1
            )

        with gr.Column(scale=1):
            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)
                
            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)

            with gr.Accordion("Model Selection", open=False):
                model_search = gr.Textbox(
                    label="Search Models",
                    placeholder="Type to filter models...",
                    lines=1
                )
                
                # 상위 9개 모델을 기본으로 설정
                default_models = [
                    "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"
                ]

                # 전체 모델 리스트
                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",
                    "Flux Animex V2",
                    "Flux Animeo V1",
                    "Flux AestheticAnime",
                    "90s Anime Art",
                    "Softserve Anime",
                    "Brain Melt Acid Art",
                    "Retro Comic Flux",
                    "Purple Dreamy",
                    "SoftPasty Flux",
                    "Flux Logo Design",
                    "Product Design",
                    "Propaganda Poster",
                    "Movie Board",
                    "Collage Flux"
                ]

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

    with gr.Row():
        generate_btn = gr.Button("Generate 2x2 Grid", variant="primary", size="lg")
        
    with gr.Row():
        gallery = gr.Gallery(
            label="Generated Images",
            show_label=True,
            elem_id="gallery",
            columns=2,
            rows=2,
            height="auto"
        )

    # 2x2 그리드 생성 함수
    def generate_grid(prompt, model, custom_lora, negative_prompt, steps, cfg_scale, seed, strength, width, height):
        images = []
        for i in range(4):
            current_seed = random.randint(1, 1000000000) if seed == -1 else seed + i
            try:
                img = query(prompt, model, custom_lora, negative_prompt, steps, cfg_scale, current_seed, strength, width, height)
                images.append(img)
            except Exception as e:
                print(f"Error generating image {i+1}: {str(e)}")
                continue
        return images

    # 이벤트 핸들러
    generate_btn.click(
        fn=generate_grid,
        inputs=[
            text_prompt,
            model,
            custom_lora,
            negative_prompt,
            steps,
            cfg,
            seed,
            strength,
            width,
            height
        ],
        outputs=gallery
    )

    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)

if __name__ == "__main__":
    dalle.launch(show_api=False, share=False)