back.py

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

# 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}')

    try:
        # 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"
            elif model == 'FLUX.1 [Dev]':
                API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
            elif model == 'FLUX.1 [Schnell]':
                API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
            elif 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}"
            elif model == 'Flux Uncensored':
                API_URL = "https://api-inference.huggingface.co/models/enhanceaiteam/Flux-uncensored"
            elif model == 'Flux Uncensored V2':
                API_URL = "https://api-inference.huggingface.co/models/enhanceaiteam/Flux-Uncensored-V2"
            elif model == 'Flux Tarot Cards':
                API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Ton618-Tarot-Cards-Flux-LoRA"
                prompt = f"Tarot card, {prompt}"
            elif model == 'Pixel Art Sprites':
                API_URL = "https://api-inference.huggingface.co/models/sWizad/pokemon-trainer-sprites-pixelart-flux"
                prompt = f"a pixel image, {prompt}"
            elif model == '3D Sketchfab':
                API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Castor-3D-Sketchfab-Flux-LoRA"
                prompt = f"3D Sketchfab, {prompt}"
            elif model == 'Retro Comic Flux':
                API_URL = "https://api-inference.huggingface.co/models/renderartist/retrocomicflux"
                prompt = f"c0m1c, comic book panel, {prompt}"
            elif model == 'Caricature':
                API_URL = "https://api-inference.huggingface.co/models/TheAwakenOne/caricature"
                prompt = f"CCTUR3, {prompt}"
            elif model == 'Huggieverse':
                API_URL = "https://api-inference.huggingface.co/models/Chunte/flux-lora-Huggieverse"
                prompt = f"HGGRE, {prompt}"
            elif model == 'Propaganda Poster':
                API_URL = "https://api-inference.huggingface.co/models/AlekseyCalvin/Propaganda_Poster_Schnell_by_doctor_diffusion"
                prompt = f"propaganda poster, {prompt}"
            elif 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}"
            elif model == 'SoftPasty Flux':
                API_URL = "https://api-inference.huggingface.co/models/alvdansen/softpasty-flux-dev"
                prompt = f"araminta_illus illustration style, {prompt}"
            elif model == 'Flux Stickers':
                API_URL = "https://api-inference.huggingface.co/models/diabolic6045/Flux_Sticker_Lora"
                prompt = f"5t1cker 5ty1e, {prompt}"
            elif model == 'Flux Animex V2':
                API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Animex-v2-LoRA"
                prompt = f"Animex, {prompt}"
            elif model == 'Flux Animeo V1':
                API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Animeo-v1-LoRA"
                prompt = f"Animeo, {prompt}"
            elif model == 'Movie Board':
                API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Flux.1-Dev-Movie-Boards-LoRA"
                prompt = f"movieboard, {prompt}"
            elif model == 'Purple Dreamy':
                API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Purple-Dreamy-Flux-LoRA"
                prompt = f"Purple Dreamy, {prompt}"
            elif model == 'PS1 Style Flux':
                API_URL = "https://api-inference.huggingface.co/models/veryVANYA/ps1-style-flux"
                prompt = f"ps1 game screenshot, {prompt}"
            elif model == 'Softserve Anime':
                API_URL = "https://api-inference.huggingface.co/models/alvdansen/softserve_anime"
                prompt = f"sftsrv style illustration, {prompt}"
            elif 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}"
            elif model == 'Half Illustration':
                API_URL = "https://api-inference.huggingface.co/models/davisbro/half_illustration"
                prompt = f"in the style of TOK, {prompt}"
            elif model == 'OpenDalle v1.1':
                API_URL = "https://api-inference.huggingface.co/models/dataautogpt3/OpenDalleV1.1"
            elif model == 'Flux Ghibsky Illustration':
                API_URL = "https://api-inference.huggingface.co/models/aleksa-codes/flux-ghibsky-illustration"
                prompt = f"GHIBSKY style, {prompt}"
            elif model == 'Flux Koda':
                API_URL = "https://api-inference.huggingface.co/models/alvdansen/flux-koda"
                prompt = f"flmft style, {prompt}"
            elif model == 'Soviet Diffusion XL':
                API_URL = "https://api-inference.huggingface.co/models/openskyml/soviet-diffusion-xl"
                prompt = f"soviet poster, {prompt}"
            elif model == 'Flux Realism LoRA':
                API_URL = "https://api-inference.huggingface.co/models/XLabs-AI/flux-RealismLora"
            elif model == 'Frosting Lane Flux':
                API_URL = "https://api-inference.huggingface.co/models/alvdansen/frosting_lane_flux"
                prompt = f"frstingln illustration, {prompt}"
            elif model == 'Phantasma Anime':
                API_URL = "https://api-inference.huggingface.co/models/alvdansen/phantasma-anime"
            elif model == 'Boreal':
                API_URL = "https://api-inference.huggingface.co/models/kudzueye/Boreal"
                prompt = f"photo, {prompt}"
            elif model == 'How2Draw':
                API_URL = "https://api-inference.huggingface.co/models/glif/how2draw"
                prompt = f"How2Draw, {prompt}"
            elif model == 'Flux AestheticAnime':
                API_URL = "https://api-inference.huggingface.co/models/dataautogpt3/FLUX-AestheticAnime"
            elif model == 'Fashion Hut Modeling LoRA':
                API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Fashion-Hut-Modeling-LoRA"
                prompt = f"Modeling of, {prompt}"
            elif model == 'Flux SyntheticAnime':
                API_URL = "https://api-inference.huggingface.co/models/dataautogpt3/FLUX-SyntheticAnime"
                prompt = f"1980s anime screengrab, VHS quality, syntheticanime, {prompt}"
            elif model == 'Flux Midjourney Anime':
                API_URL = "https://api-inference.huggingface.co/models/brushpenbob/flux-midjourney-anime"
                prompt = f"egmid, {prompt}"
            elif model == 'Coloring Book Generator':
                API_URL = "https://api-inference.huggingface.co/models/robert123231/coloringbookgenerator"
            elif model == 'Collage Flux':
                API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Castor-Collage-Dim-Flux-LoRA"
                prompt = f"collage, {prompt}"
            elif model == 'Flux Product Ad Backdrop':
                API_URL = "https://api-inference.huggingface.co/models/prithivMLmods/Flux-Product-Ad-Backdrop"
                prompt = f"Product Ad, {prompt}"
            elif model == 'Product Design':
                API_URL = "https://api-inference.huggingface.co/models/multimodalart/product-design"
                prompt = f"product designed by prdsgn, {prompt}"
            elif model == '90s Anime Art':
                API_URL = "https://api-inference.huggingface.co/models/glif/90s-anime-art"
            elif 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}"
            elif model == 'Lustly Flux Uncensored v1':
                API_URL = "https://api-inference.huggingface.co/models/lustlyai/Flux_Lustly.ai_Uncensored_nsfw_v1"
            elif model == 'NSFW Master Flux':
                API_URL = "https://api-inference.huggingface.co/models/Keltezaa/NSFW_MASTER_FLUX"
                prompt = f"NSFW, {prompt}"
            elif model == 'Flux Outfit Generator':
                API_URL = "https://api-inference.huggingface.co/models/tryonlabs/FLUX.1-dev-LoRA-Outfit-Generator"
            elif model == 'Midjourney':
                API_URL = "https://api-inference.huggingface.co/models/Jovie/Midjourney"
            elif model == 'DreamPhotoGASM':
                API_URL = "https://api-inference.huggingface.co/models/Yntec/DreamPhotoGASM"
            elif model == 'Flux Super Realism LoRA':
                API_URL = "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Super-Realism-LoRA"
            elif model == 'Stable Diffusion 2-1':
                API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2-1-base"
            elif model == 'Stable Diffusion 3.5 Large':
                API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large"
            elif model == 'Stable Diffusion 3.5 Large Turbo':
                API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large-turbo"
            elif model == 'Stable Diffusion 3 Medium':
                API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3-medium-diffusers"
                prompt = f"A, {prompt}"
            elif model == 'Duchaiten Real3D NSFW XL':
                API_URL = "https://api-inference.huggingface.co/models/stablediffusionapi/duchaiten-real3d-nsfw-xl"
            elif model == 'Pixel Art XL':
                API_URL = "https://api-inference.huggingface.co/models/nerijs/pixel-art-xl"
                prompt = f"pixel art, {prompt}"
            elif model == 'Character Design':
                API_URL = "https://api-inference.huggingface.co/models/KappaNeuro/character-design"
                prompt = f"Character Design, {prompt}"
            elif model == 'Sketched Out Manga':
                API_URL = "https://api-inference.huggingface.co/models/alvdansen/sketchedoutmanga"
                prompt = f"daiton, {prompt}"
            elif model == 'Archfey Anime':
                API_URL = "https://api-inference.huggingface.co/models/alvdansen/archfey_anime"
            elif model == 'Lofi Cuties':
                API_URL = "https://api-inference.huggingface.co/models/alvdansen/lofi-cuties"
            elif model == 'YiffyMix':
                API_URL = "https://api-inference.huggingface.co/models/Yntec/YiffyMix"
            elif model == 'Analog Madness Realistic v7':
                API_URL = "https://api-inference.huggingface.co/models/digiplay/AnalogMadness-realistic-model-v7"
            elif 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}"
            elif model == 'Filmgrain':
                API_URL = "https://api-inference.huggingface.co/models/artificialguybr/filmgrain-redmond-filmgrain-lora-for-sdxl"
                prompt = f"Film Grain, FilmGrainAF, {prompt}"
            elif 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}"
            elif model == 'Cyborg Style XL':
                API_URL = "https://api-inference.huggingface.co/models/goofyai/cyborg_style_xl"
                prompt = f"cyborg style, {prompt}"
            elif model == 'Little Tinies':
                API_URL = "https://api-inference.huggingface.co/models/alvdansen/littletinies"
            elif model == 'NSFW XL':
                API_URL = "https://api-inference.huggingface.co/models/Dremmar/nsfw-xl"
            elif model == 'Analog Redmond':
                API_URL = "https://api-inference.huggingface.co/models/artificialguybr/analogredmond"
                prompt = f"timeless style, {prompt}"
            elif model == 'Pixel Art Redmond':
                API_URL = "https://api-inference.huggingface.co/models/artificialguybr/PixelArtRedmond"
                prompt = f"Pixel Art, {prompt}"
            elif model == 'Ascii Art':
                API_URL = "https://api-inference.huggingface.co/models/CiroN2022/ascii-art"
                prompt = f"ascii art, {prompt}"
            elif model == 'Analog':
                API_URL = "https://api-inference.huggingface.co/models/Yntec/Analog"
            elif model == 'Maple Syrup':
                API_URL = "https://api-inference.huggingface.co/models/Yntec/MapleSyrup"
            elif model == 'Perfect Lewd Fantasy':
                API_URL = "https://api-inference.huggingface.co/models/digiplay/perfectLewdFantasy_v1.01"
            elif model == 'AbsoluteReality 1.8.1':
                API_URL = "https://api-inference.huggingface.co/models/digiplay/AbsoluteReality_v1.8.1"
            elif model == 'Disney':
                API_URL = "https://api-inference.huggingface.co/models/goofyai/disney_style_xl"
                prompt = f"Disney style, {prompt}"
            elif model == 'Redmond SDXL':
                API_URL = "https://api-inference.huggingface.co/models/artificialguybr/LogoRedmond-LogoLoraForSDXL-V2"
            elif model == 'epiCPhotoGasm':
                API_URL = "https://api-inference.huggingface.co/models/Yntec/epiCPhotoGasm"
            else:
                API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"

        # 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
            }
        }

        # 타임아웃 값을 늘리고 재시도 로직 추가
        max_retries = 3
        current_retry = 0
        
        while current_retry < max_retries:
            try:
                response = requests.post(API_URL, headers=headers, json=payload, timeout=180)  # 타임아웃을 180초로 증가
                response.raise_for_status()
                
                image = Image.open(io.BytesIO(response.content))
                print(f'Generation {key} completed successfully')
                return image
                
            except requests.exceptions.Timeout:
                current_retry += 1
                if current_retry < max_retries:
                    print(f"Timeout occurred. Retrying... (Attempt {current_retry + 1}/{max_retries})")
                    continue
                else:
                    raise gr.Error(f"Request timed out after {max_retries} attempts. The model might be busy, please try again later.")
                    
            except requests.exceptions.RequestException as e:
                raise gr.Error(f"Request failed: {str(e)}")
        
    except requests.exceptions.RequestException as e:
        error_message = f"Request failed: {str(e)}"
        if hasattr(e, 'response') and e.response is not None:
            if e.response.status_code == 401:
                error_message = "Invalid API token. Please check your Hugging Face API token."
            elif e.response.status_code == 403:
                error_message = "Access denied. Please check your API token permissions."
            elif e.response.status_code == 503:
                error_message = "Model is currently loading. Please try again in a few moments."
        raise gr.Error(error_message)
    except Exception as e:
        raise gr.Error(f"Unexpected error: {str(e)}")


def generate_grid(prompt, selected_models, custom_lora, negative_prompt, steps, cfg_scale, seed, strength, width, height, progress=gr.Progress()):
    if len(selected_models) > 4:
        raise gr.Error("Please select up to 4 models")
    if len(selected_models) == 0:
        raise gr.Error("Please select at least 1 model")
    
    # 초기 이미지 배열 생성
    images = []
    total_models = len(selected_models[:4])
    
    # 각 모델별로 이미지 생성
    for idx, model_name in enumerate(selected_models[:4]):
        try:
            progress((idx + 1) / total_models, f"Generating image for {model_name}...")
            img = query(prompt, model_name, custom_lora, negative_prompt, steps, cfg_scale, seed, strength, width, height)
            if img:
                images.append(img)
        except Exception as e:
            print(f"Error generating image for {model_name}: {str(e)}")
            continue
    
    # 최소한 하나의 이미지가 생성되었는지 확인
    if not images:
        raise gr.Error("Failed to generate any images. Please try again.")
    
    # 4개의 이미지 슬롯을 채움
    while len(images) < 4:
        images.append(images[-1])
    
    # 이미지가 올바르게 로드되었는지 확인
    valid_images = []
    for img in images:
        try:
            # 이미지 복사본 생성
            img_copy = img.copy()
            valid_images.append(img_copy)
        except Exception as e:
            print(f"Error processing image: {str(e)}")
            # 오류가 발생한 경우 마지막 유효한 이미지로 대체
            if valid_images:
                valid_images.append(valid_images[-1].copy())
            else:
                raise gr.Error("Failed to process images. Please try again.")
    
    progress(1.0, "Generation complete!")
    return valid_images

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

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
    gr.Markdown("# ZeroWeight Studio")
    
    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
                )

                # 상위 4개 모델을 기본으로 설정
                default_models = [
                    "FLUX.1 [Schnell]",  # 모델 이름 통일
                    "Stable Diffusion 3.5 Large",
                    "Stable Diffusion 3.5 Large Turbo",
                    "Midjourney"
                ]

                # 전체 모델 리스트
                models_list = [
                    "FLUX.1 [Schnell]",  # 모델 이름 통일
                    "Stable Diffusion 3.5 Large",
                    "Stable Diffusion 3.5 Large Turbo",
                    "Stable Diffusion XL",
                    "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.Checkboxgroup(
                    label="Select Models (Choose up to 4)",
                    choices=models_list,
                    value=default_models,
                    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"
        )

    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, value=[])

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

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