app.py

import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import DiffusionPipeline, AutoencoderTiny
from huggingface_hub import hf_hub_download

import config
styles_name = [style["name"] for style in config.style_list]
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2240

def feifeimodload():

    dtype = torch.bfloat16
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    pipe = DiffusionPipeline.from_pretrained(
        "aifeifei798/DarkIdol-flux-v1.1", torch_dtype=dtype
    ).to(device)

    #pipe = DiffusionPipeline.from_pretrained(
    #    "city96/Flux.1-Heavy-17B", torch_dtype=dtype
    #).to(device)
    
    pipe.vae.enable_slicing()
    pipe.vae.enable_tiling()
    #pipe.unload_lora_weights()
    torch.cuda.empty_cache()
    return pipe
    
pipe = feifeimodload()

@spaces.GPU()
def infer(prompt, styles_Radio="(None)", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    if styles_Radio:
        for style in config.style_list:
            if style["name"] == styles_Radio:
                prompt = style["prompt"].replace("{prompt}", prompt)
    image = pipe(
            prompt = "flux, 8k, ",
            prompt_2 = prompt,  
            width = width,
            height = height,
            num_inference_steps = num_inference_steps, 
            generator = generator,
            guidance_scale=3.5
    ).images[0] 
    return image, seed
 
examples = [
    "A group of girls in bikinis playing orchestra",
    "this photo is a smile girl in cute bikini",
    "this photo is a cute girl in cute bikini",
    "kpop girls, sunrise",
    "If life could always be as fresh as the first encounter.",
    "Dark Idol this photo is a sexy cute girl in cute bikini",
    "a sexy girl,poses,look at camera,Slim figure, gigantic breasts,poses,natural,High-quality photography, creative composition, fashion foresight, a strong visual style, and an aura of luxury and sophistication collectively define the distinctive aesthetic of Vogue magazine.",
    "real model slight smile girl in real life",
    "real model smile girl in real life",
    "real model girl in real life",
    "A high-resolution photograph of a Japanese female model in a serene, natural setting, with soft, warm lighting, and a minimalist aesthetic, showcasing a elegant fragrance bottle and the model's effortless, emotive expression, with impeccable styling, and a muted color palette, evoking a sense of understated luxury and refinement."
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# DarkIdol-flux
DarkIdol-flux is a text-to-image AI model designed to create aesthetic, detailed and diverse images from textual prompts in just 6-8 steps. It offers enhanced performance in image quality, typography, understanding complex prompts, and resource efficiency.
        """)
        
        with gr.Row():
            
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=12,
                placeholder="Enter your prompt",
                container=False,
                value="",
            )
        with gr.Row():
            styles_Radio = gr.Dropdown(
                styles_name,
                label="Styles",
                multiselect=False,
                value="(None)",
            )
        run_button = gr.Button("Run")
        
        result = gr.Image(label="Result", show_label=False,height=520)
        
        with gr.Accordion("Advanced Settings", open=False):
            
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            
            with gr.Row():
                
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1088,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1920,
                )
            
            with gr.Row():
                
  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4,
                )
        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [result, seed],
            cache_examples=False
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs = [prompt, styles_Radio, seed, randomize_seed, width, height, num_inference_steps],
        outputs = [result, seed]
    )

demo.launch()