app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces
import re
import os
from diffusers import (
    DiffusionPipeline,
    AutoencoderTiny,
)
from huggingface_hub import hf_hub_download
#from feifeilib.feifeichat import feifeichat

IS_ZERO_GPU = bool(os.getenv("SPACES_ZERO_GPU"))
IS_GPU_MODE = True if IS_ZERO_GPU else (True if torch.cuda.is_available() else False)
if IS_ZERO_GPU:
    import subprocess
    subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
    torch.set_float32_matmul_precision("high")
    torch.backends.cuda.matmul.allow_tf32 = True

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", torch_dtype=dtype
    ).to(device)
    
    pipe.load_lora_weights(
        hf_hub_download("aifeifei798/feifei-flux-lora-v1.1", "feifei-v1.1.safetensors"),
        adapter_name="feifei",
    )
    pipe.load_lora_weights(
        hf_hub_download("aifeifei798/sldr_flux_nsfw_v2-studio", "sldr_flux_nsfw_v2-studio.safetensors"),
        adapter_name="sldr_flux_nsfw_v2",
    )

    pipe.load_lora_weights(
        hf_hub_download(
            "aifeifei798/big-boobs-clothed",
            "big-boobs-clothed-v2.safetensors",
        ),
        adapter_name="big-boobs-clothed-v2",
    )    
    
    pipe.vae.enable_slicing()
    pipe.vae.enable_tiling()
    torch.cuda.empty_cache()
    return pipe
    
pipe = feifeimodload()

if IS_ZERO_GPU:
    from optimization import optimize_pipeline_
    optimize_pipeline_(pipe, "prompt")

def get_duration(prompt, styles_Radio, feife_select, bigboobs_select, seed, randomize_seed, width, height, num_inference_steps, guidancescale, num_feifei, nsfw_select, nsfw_slider, progress):
    def_duration = 30.
    def_steps = 4.
    return int(def_duration * float(num_inference_steps) / def_steps)

@spaces.GPU(duration=get_duration)
def infer(prompt="",  styles_Radio="(None)", feife_select = False, bigboobs_select = True, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidancescale=3.5, num_feifei=0.35, nsfw_select=False, nsfw_slider=1, progress=gr.Progress(track_tqdm=True)):
    
    Duke86Syl_lora_name=[]
    adapter_weights_num=[]
    
    if feife_select:
        Duke86Syl_lora_name.append("feifei")
        adapter_weights_num.append(num_feifei)
        
    if bigboobs_select:
        Duke86Syl_lora_name.append("big-boobs-clothed-v2")
        adapter_weights_num.append(0.45)

    if nsfw_select:
        Duke86Syl_lora_name.append("sldr_flux_nsfw_v2")
        adapter_weights_num.append(nsfw_slider)

    pipe.set_adapters(
        Duke86Syl_lora_name,
        adapter_weights=adapter_weights_num,
    )
    pipe.fuse_lora(
        adapter_name=Duke86Syl_lora_name,
        lora_scale=1.0,
    )
        
    #if feife_select:
    #    # Define the replacement string
    #    replacement = " feifei, A beautiful, 18 yo kpop idol, large-busted Japanese slim girl, with light makeup, gazing deeply into the camera, "
    #    
    #    # Perform the replacements with re.IGNORECASE
    #    prompt = re.sub(r"girl", replacement, prompt, flags=re.IGNORECASE)
    #    prompt = re.sub(r"young woman", replacement, prompt, flags=re.IGNORECASE)
    #    prompt = re.sub(r"woman", replacement, prompt, flags=re.IGNORECASE)
    #    prompt = re.sub(r"model", replacement, prompt, flags=re.IGNORECASE)
        
    if styles_Radio:
        style_name = styles_Radio
        for style in config.style_list:
            if style["name"] == style_name:
                prompt = style["prompt"].replace("{prompt}", prompt)
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
       
    image = pipe(
            prompt = "flux, 8k, ",
            prompt_2 = prompt,
            width = width,
            height = height,
            num_inference_steps = num_inference_steps, 
            generator = generator,
            guidance_scale=guidancescale
    ).images[0] 
    return image, seed
 
examples = [
    "this photo is a girl",
    "this photo is a girl in bikini",
    "this photo is a cute girl in cute bikini",
    "girl, sunrise",
    "DarkIdol flux girl",
    "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 girl 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"):
        result = gr.Image(label="Result", show_label=False,height=500,format="png")
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=12,
                placeholder="Enter your prompt",
                value="",
                container=False,
            )
        run_button = gr.Button("Run")
        with gr.Row():
            styles_Radio = gr.Dropdown(
                styles_name,
                label="Styles",
                multiselect=False,
                value="(None)",
            )
        #feifei_chat_text1 = gr.Textbox(
        #    label="输入您的问题1：", show_label=False, container=False, lines=1, value="")
        #feifei_chat_text2 = gr.Textbox(
        #    label="输入您的问题2：", show_label=False, container=False, lines=3, value="")
        
        # 定义模型选择下拉框
        #feifei_chat_Dropdown = gr.Dropdown(
        #    [
        #        "meta/llama-3.3-70b-instruct",
        #        "nvidia/llama-3.3-nemotron-super-49b-v1",
        #        "mistralai/Mistral-Nemo-Instruct-2411",
        #    ],
        #    value="mistralai/Mistral-Nemo-Instruct-2411",
        #    label="选择模型", show_label=False, container=False
        #)
        # 定义提交按钮
        #feifei_chat_btn = gr.Button(value="Gen Prompt")
        feife_select = gr.Checkbox(label="FeiFei Expansion", value=True) 
        bigboobs_select = gr.Checkbox(label="bigboobs", value=True) 
        nsfw_select = gr.Checkbox(label="NSFW") 
        nsfw_slider = gr.Slider(
                    label="NSFW",
                    minimum=0,
                    maximum=2,
                    step=0.05,
                    value=0.75,
                )
        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():
                gr.Markdown(''' - 21:9 2240x1024
                            - 16:9 1856x1024
                            - 9:7  1344x1024 ''')                
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1024,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1856,
                )
            
            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4,
                )

            with gr.Row():
                guidancescale = gr.Slider(
                    label="Guidance scale",
                    minimum=0,
                    maximum=10,
                    step=0.1,
                    value=3.5,
                )
                
            with gr.Row():
                num_feifei = gr.Slider(
                    label="FeiFei",
                    minimum=0,
                    maximum=2,
                    step=0.05,
                    value=0.35,
                )

        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [result, seed],
            cache_examples=False
        )

    run_button.click(
        fn = infer,
        inputs = [prompt, styles_Radio, feife_select, bigboobs_select, seed, randomize_seed, width, height, num_inference_steps, guidancescale, num_feifei, nsfw_select, nsfw_slider],
        outputs = [result, seed]
    )
    #feifei_chat_btn.click(
    #    fn=feifeichat,
    #    inputs=[feifei_chat_text1, feifei_chat_text2, feifei_chat_Dropdown],
    #    outputs=prompt
    #)

demo.launch()