Update app.py

app.py

@@ -10,6 +10,8 @@ from diffusers import (
 )
 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 = 2048
 
@@ -26,7 +28,10 @@ def feifeimodload():
         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.vae.enable_slicing()
     pipe.vae.enable_tiling()
     torch.cuda.empty_cache()
@@ -35,7 +40,42 @@ def feifeimodload():
 pipe = feifeimodload()
 
 @spaces.GPU()
-def infer(prompt="", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidancescale=3.5, num_feifei=0.35, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt="",  styles_Radio="(None)", Expansion_select=False, 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)):
+    
+    if nsfw_select:
+        pipe.set_adapters(
+            ["feifei", "sldr_flux_nsfw_v2"],
+            adapter_weights=[num_feifei,nsfw_slider],
+        )
+        pipe.fuse_lora(
+            adapter_name=["feifei", "sldr_flux_nsfw_v2"],
+            lora_scale=1.0,
+        )
+    else:
+        pipe.set_adapters(
+            ["feifei"],
+            adapter_weights=[num_feifei],
+        )
+        pipe.fuse_lora(
+            adapter_name=["feifei"],
+            lora_scale=1.0,
+        )
+        
+    if Expansion_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)
@@ -94,8 +134,15 @@ with gr.Blocks(css=css) as demo:
                 container=False,
             )
         run_button = gr.Button("Run")
+        with gr.Row():
+            styles_Radio = gr.Dropdown(
+                styles_name,
+                label="Styles",
+                multiselect=False,
+                value="(None)",
+            )
         result = gr.Image(label="Result", show_label=False,height=520)
-        
+        Expansion_select = gr.Checkbox(label="FeiFei Expansion") 
         with gr.Accordion("Advanced Settings", open=False):
             
             seed = gr.Slider(
@@ -152,7 +199,15 @@ with gr.Blocks(css=css) as demo:
                     step=0.05,
                     value=0.45,
                 )
-
+            with gr.Row():
+                nsfw_select = gr.Checkbox(label="NSFW") 
+                nsfw_slider = gr.Slider(
+                    label="NSFW",
+                    minimum=0,
+                    maximum=2,
+                    step=0.05,
+                    value=0.75,
+                )
         
         gr.Examples(
             examples = examples,
@@ -165,7 +220,7 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = infer,
-        inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps, guidancescale, num_feifei],
+        inputs = [prompt, styles_Radio, Expansion_select, seed, randomize_seed, width, height, num_inference_steps, guidancescale, num_feifei, nsfw_select, nsfw_slider],
         outputs = [result, seed]
     )