Upload app.py

app.py

@@ -1,33 +1,274 @@
-from datetime import datetime
-
 import gradio as gr
-import spaces
+import numpy as np
+import random
 import torch
-from diffusers import FluxPipeline
+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",
+    )
 
-from optimization import optimize_pipeline_
+    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")
 
-pipeline = FluxPipeline.from_pretrained('black-forest-labs/FLUX.1-dev', torch_dtype=torch.bfloat16).to('cuda')
-optimize_pipeline_(pipeline, "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)
 
-@spaces.GPU
-def generate_image(prompt: str, progress=gr.Progress(track_tqdm=True)):
-    generator = torch.Generator(device='cuda').manual_seed(42)
-    t0 = datetime.now()
-    output = pipeline(
-        prompt=prompt,
-        num_inference_steps=28,
-        generator=generator,
+    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,
     )
-    return [(output.images[0], f'{(datetime.now() - t0).total_seconds():.2f}s')]
+        
+    #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
+    #)
 
-gr.Interface(
-    fn=generate_image,
-    inputs=gr.Text(label="Prompt"),
-    outputs=gr.Gallery(),
-    examples=["A cat playing with a ball of yarn"],
-    cache_examples=False,
-).launch()
+demo.launch()