Upload 4 files

app.py

@@ -8,10 +8,12 @@ from diffusers import FluxControlNetPipeline, FluxControlNetModel, FluxMultiCont
 from huggingface_hub import HfFileSystem, ModelCard
 import random
 import time
+import os
 
-from mod import (models, clear_cache, get_repo_safetensors, is_repo_name, is_repo_exists,
-                 description_ui, num_loras, compose_lora_json, is_valid_lora, fuse_loras,
-                 get_trigger_word, enhance_prompt, deselect_lora, num_cns, set_control_union_image,
+from env import models, num_loras, num_cns
+from mod import (clear_cache, get_repo_safetensors, is_repo_name, is_repo_exists, get_model_trigger,
+                 description_ui, compose_lora_json, is_valid_lora, fuse_loras, save_image,
+                 get_trigger_word, enhance_prompt, deselect_lora, set_control_union_image,
                  get_control_union_mode, set_control_union_mode, get_control_params, translate_to_en)
 from flux import (search_civitai_lora, select_civitai_lora, search_civitai_lora_json,
                   download_my_lora, get_all_lora_tupled_list, apply_lora_prompt,
@@ -26,7 +28,7 @@ with open('loras.json', 'r') as f:
 # Initialize the base model
 base_model = models[0]
 controlnet_model_union_repo = 'InstantX/FLUX.1-dev-Controlnet-Union'
-#controlnet_model_union_repo = 'InstantX/FLUX.1-dev-Controlnet-Union-alpha'
+#controlnet_model_union_repo = 'Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro'
 dtype = torch.bfloat16
 #dtype = torch.float8_e4m3fn
 #device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -128,6 +130,7 @@ def update_selection(evt: gr.SelectData, width, height):
 
 @spaces.GPU(duration=70)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, cn_on, progress=gr.Progress(track_tqdm=True)):
+    from diffusers.utils import load_image
     global pipe
     global taef1
     global good_vae
@@ -142,8 +145,8 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
             # Generate image
             modes, images, scales = get_control_params()
             if not cn_on or len(modes) == 0:
-                pipe.vae = taef1
                 pipe.to("cuda")
+                pipe.vae = taef1
                 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
                 progress(0, desc="Start Inference.")
                 for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
@@ -159,11 +162,11 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
                 ):
                     yield img
             else:
+                pipe.to("cuda")
                 if controlnet is not None: controlnet.to("cuda")
                 if controlnet_union is not None: controlnet_union.to("cuda")
                 pipe.vae = good_vae
                 pipe.controlnet = controlnet
-                pipe.to("cuda")
                 progress(0, desc="Start Inference with ControlNet.")
                 for img in pipe(
                     prompt=prompt_mash,
@@ -201,7 +204,7 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
 
     if translate_on: prompt = translate_to_en(prompt)
 
-    prompt_mash = prompt
+    prompt_mash = prompt + get_model_trigger(last_model)
     if is_valid_lora(lora_json):
         # Load External LoRA weights
         with calculateDuration("Loading External LoRA weights"):
@@ -215,13 +218,13 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
         if(trigger_word):
             if "trigger_position" in selected_lora:
                 if selected_lora["trigger_position"] == "prepend":
-                    prompt_mash = f"{trigger_word} {prompt}"
+                    prompt_mash = f"{trigger_word} {prompt_mash}"
                 else:
-                    prompt_mash = f"{prompt} {trigger_word}"
+                    prompt_mash = f"{prompt_mash} {trigger_word}"
             else:
-                prompt_mash = f"{trigger_word} {prompt}"
+                prompt_mash = f"{trigger_word} {prompt_mash}"
         else:
-            prompt_mash = prompt
+            prompt_mash = prompt_mash
         # Load LoRA weights
         with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
             if "weights" in selected_lora:
@@ -245,7 +248,7 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
         progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
         yield image, seed, gr.update(value=progress_bar, visible=True)
         
-    yield final_image, seed, gr.update(value=progress_bar, visible=False)
+    yield save_image(final_image, None, last_model, prompt_mash, height, width, steps, cfg_scale, seed), seed, gr.update(value=progress_bar, visible=False)
 
 def get_huggingface_safetensors(link):
   split_link = link.split("/")
@@ -338,7 +341,7 @@ css = '''
 #progress .generating{display:none}
 .progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px}
 .progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
-.info {text-align:center; display:inline-flex; align-items:center !important}
+.info {text-align:center; !important}
 '''
 with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css, delete_cache=(60, 3600)) as app:
     with gr.Tab("FLUX LoRA the Explorer"):
@@ -460,7 +463,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css, delete_cache
                                         with gr.Row():
                                             cn_image_ref[i] = gr.Image(label="Image Reference", type="pil", format="png", height=256, sources=["upload", "clipboard"], show_share_button=False)
                                             cn_image[i] = gr.Image(label="Control Image", type="pil", format="png", height=256, show_share_button=False, interactive=False)
-
+    
     gallery.select(
         update_selection,
         inputs=[width, height],
@@ -643,84 +646,84 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css, delete_cache
                 pg_generate_text_button = gr.Button("Generate Prompt with LLM (Llama 3.1 70B)")
                 pg_text_output = gr.Textbox(label="Generated Text", lines=10)
 
-    description_ui()
-
-    def create_caption(image):
-        if image is not None:
-            return florence_caption(image)
-        return ""
+        def create_caption(image):
+            if image is not None:
+                return florence_caption(image)
+            return ""
 
-    pg_create_caption_button.click(
-        create_caption,
-        inputs=[pg_input_image],
-        outputs=[pg_caption_output]
-    )
+        pg_create_caption_button.click(
+            create_caption,
+            inputs=[pg_input_image],
+            outputs=[pg_caption_output]
+        )
 
-    def generate_prompt_with_dynamic_seed(*args):
-        # Generate a new random seed
-        dynamic_seed = random.randint(0, 1000000)
-        
-        # Call the generate_prompt function with the dynamic seed
-        result = prompt_generator.generate_prompt(dynamic_seed, *args)
-        
-        # Return the result along with the used seed
-        return [dynamic_seed] + list(result)
-
-    pg_generate_button.click(
-        generate_prompt_with_dynamic_seed,
-        inputs=[pg_custom, pg_subject, pg_gender, pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles,
-                pg_additional_details, pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform,
-                pg_place, pg_lighting, pg_clothing, pg_composition, pg_pose, pg_background, pg_input_image],
-        outputs=[gr.Number(label="Used Seed", visible=False), pg_output, gr.Number(visible=False), pg_t5xxl_output, pg_clip_l_output, pg_clip_g_output]
-    ) #
-
-    pg_add_caption_button.click(
-        prompt_generator.add_caption_to_prompt,
-        inputs=[pg_output, pg_caption_output],
-        outputs=[pg_output]
-    )
+        def generate_prompt_with_dynamic_seed(*args):
+            # Generate a new random seed
+            dynamic_seed = random.randint(0, 1000000)
+            
+            # Call the generate_prompt function with the dynamic seed
+            result = prompt_generator.generate_prompt(dynamic_seed, *args)
+            
+            # Return the result along with the used seed
+            return [dynamic_seed] + list(result)
+
+        pg_generate_button.click(
+            generate_prompt_with_dynamic_seed,
+            inputs=[pg_custom, pg_subject, pg_gender, pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles,
+                    pg_additional_details, pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform,
+                    pg_place, pg_lighting, pg_clothing, pg_composition, pg_pose, pg_background, pg_input_image],
+            outputs=[gr.Number(label="Used Seed", visible=False), pg_output, gr.Number(visible=False), pg_t5xxl_output, pg_clip_l_output, pg_clip_g_output]
+        ) #
+
+        pg_add_caption_button.click(
+            prompt_generator.add_caption_to_prompt,
+            inputs=[pg_output, pg_caption_output],
+            outputs=[pg_output]
+        )
 
-    pg_generate_text_button.click(
-        huggingface_node.generate,
-        inputs=[pg_output, pg_happy_talk, pg_compress, pg_compression_level, pg_poster, pg_custom_base_prompt],
-        outputs=pg_text_output
-    )
+        pg_generate_text_button.click(
+            huggingface_node.generate,
+            inputs=[pg_output, pg_happy_talk, pg_compress, pg_compression_level, pg_poster, pg_custom_base_prompt],
+            outputs=pg_text_output
+        )
 
-    def update_all_options(choice):
-        updates = {}
-        if choice == "Disabled":
-            for dropdown in [
-                pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing,
-                pg_place, pg_lighting, pg_composition, pg_pose, pg_background, pg_additional_details,
-                pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform
-            ]:
-                updates[dropdown] = gr.update(value="disabled")
-        elif choice == "Random":
-            for dropdown in [
+        def update_all_options(choice):
+            updates = {}
+            if choice == "Disabled":
+                for dropdown in [
+                    pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing,
+                    pg_place, pg_lighting, pg_composition, pg_pose, pg_background, pg_additional_details,
+                    pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform
+                ]:
+                    updates[dropdown] = gr.update(value="disabled")
+            elif choice == "Random":
+                for dropdown in [
+                    pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing,
+                    pg_place, pg_lighting, pg_composition, pg_pose, pg_background, pg_additional_details,
+                    pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform
+                ]:
+                    updates[dropdown] = gr.update(value="random")
+            else:  # No Figure Random
+                for dropdown in [pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing, pg_pose, pg_additional_details]:
+                    updates[dropdown] = gr.update(value="disabled")
+                for dropdown in [pg_artform, pg_place, pg_lighting, pg_composition, pg_background, pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform]:
+                    updates[dropdown] = gr.update(value="random")
+            return updates
+        
+        pg_global_option.change(
+            update_all_options,
+            inputs=[pg_global_option],
+            outputs=[
                 pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing,
                 pg_place, pg_lighting, pg_composition, pg_pose, pg_background, pg_additional_details,
                 pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform
-            ]:
-                updates[dropdown] = gr.update(value="random")
-        else:  # No Figure Random
-            for dropdown in [pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing, pg_pose, pg_additional_details]:
-                updates[dropdown] = gr.update(value="disabled")
-            for dropdown in [pg_artform, pg_place, pg_lighting, pg_composition, pg_background, pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform]:
-                updates[dropdown] = gr.update(value="random")
-        return updates
-    
-    pg_global_option.change(
-        update_all_options,
-        inputs=[pg_global_option],
-        outputs=[
-            pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing,
-            pg_place, pg_lighting, pg_composition, pg_pose, pg_background, pg_additional_details,
-            pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform
-        ]
-    )
+            ]
+        )
 
+    description_ui()
     gr.LoginButton()
     gr.DuplicateButton(value="Duplicate Space for private use (This demo does not work on CPU. Requires GPU Space)")
 
+
 app.queue()
 app.launch()

env.py

@@ -6,6 +6,61 @@ hf_token = os.environ.get("HF_TOKEN")
 hf_read_token = os.environ.get('HF_READ_TOKEN') # only use for private repo
 
 
+num_loras = 3
+num_cns = 2
+
+
+models = [
+    "camenduru/FLUX.1-dev-diffusers",
+    "black-forest-labs/FLUX.1-schnell",
+    "sayakpaul/FLUX.1-merged",
+    "ostris/OpenFLUX.1",
+    "John6666/flux-dev2pro-bf16-flux",
+    "John6666/flux1-dev-minus-v1-fp8-flux",
+    "John6666/hyper-flux1-dev-fp8-flux",
+    "John6666/blue-pencil-flux1-v010-fp8-flux",
+    "Raelina/Raemu-Flux",
+    "John6666/raemu-flux-v10-fp8-flux",
+    "John6666/copycat-flux-test-fp8-v11-fp8-flux",
+    "John6666/wai-ani-flux-v10forfp8-fp8-flux",
+    "John6666/flux-dev8-anime-nsfw-fp8-flux",
+    "John6666/nepotism-fuxdevschnell-v3aio-fp8-flux",
+    "John6666/sumeshi-flux1s-v002e-fp8-flux",
+    "John6666/fca-style-v33-x10-8step-fp8-flux",
+    "John6666/lyh-anime-v10f1-fp8-flux",
+    "John6666/lyh-dalle-anime-v12dalle-fp8-flux",
+    "John6666/lyh-anime-flux-v2a1-fp8-flux",
+    "John6666/glimmerkin-flux-cute-v10-fp8-flux",
+    "John6666/niji-style-flux-devfp8-fp8-flux",
+    "John6666/niji56-style-v3-fp8-flux",
+    "John6666/xe-anime-flux-v04-fp8-flux",
+    "John6666/xe-figure-flux-01-fp8-flux",
+    "John6666/xe-pixel-flux-01-fp8-flux",
+    "John6666/xe-guoman-flux-02-fp8-flux",
+    "John6666/carnival-unchained-v10-fp8-flux",
+    "John6666/fluxunchained-artfulnsfw-fut516xfp8e4m3fnv11-fp8-flux",
+    "John6666/fastflux-unchained-t5f16-fp8-flux",
+    "John6666/iniverse-mix-xl-sfwnsfw-fluxdfp16nsfwv11-fp8-flux",
+    "John6666/nsfw-master-flux-lora-merged-with-flux1-dev-fp16-v10-fp8-flux",
+    "John6666/the-araminta-flux1a1-fp8-flux",
+    "John6666/acorn-is-spinning-flux-v11-fp8-flux",
+    "John6666/real-horny-pro-fp8-flux",
+    "John6666/centerfold-flux-v20fp8e5m2-fp8-flux",
+    "John6666/jib-mix-flux-v208stephyper-fp8-flux",
+    "John6666/sapianf-nude-men-women-for-flux-v20fp16-fp8-flux",
+    "John6666/flux-asian-realistic-v10-fp8-flux",
+    "John6666/fluxasiandoll-v10-fp8-flux",
+    "John6666/xe-asian-flux-01-fp8-flux",
+    "John6666/fluxescore-dev-v10fp16-fp8-flux",
+    # "",
+]
+
+model_trigger = {
+    "Raelina/Raemu-Flux": "anime",
+    "John6666/raemu-flux-v10-fp8-flux": "anime",
+    "John6666/fca-style-v33-x10-8step-fp8-flux": "fca_style",
+}
+
 # List all Models for specified user
 HF_MODEL_USER_LIKES = [] # sorted by number of likes
 HF_MODEL_USER_EX = [] # sorted by a special rule
@@ -37,4 +92,3 @@ HF_LORA_PRIVATE_REPOS2 = [] # to be sorted as 1 repo
 HF_LORA_PRIVATE_REPOS = HF_LORA_PRIVATE_REPOS1 + HF_LORA_PRIVATE_REPOS2
 HF_LORA_ESSENTIAL_PRIVATE_REPO = '' # to be downloaded on run app
 HF_VAE_PRIVATE_REPO = ''
-

mod.py

@@ -5,6 +5,7 @@ from PIL import Image
 from pathlib import Path
 import gc
 import subprocess
+from env import num_cns, model_trigger
 
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
@@ -13,54 +14,6 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 torch.set_grad_enabled(False)
 
 
-models = [
-    "camenduru/FLUX.1-dev-diffusers",
-    "black-forest-labs/FLUX.1-schnell",
-    "sayakpaul/FLUX.1-merged",
-    "ostris/OpenFLUX.1",
-    "John6666/flux-dev2pro-bf16-flux",
-    "John6666/flux1-dev-minus-v1-fp8-flux",
-    "John6666/hyper-flux1-dev-fp8-flux",
-    "John6666/blue-pencil-flux1-v010-fp8-flux",
-    "Raelina/Raemu-Flux",
-    "John6666/raemu-flux-v10-fp8-flux",
-    "John6666/copycat-flux-test-fp8-v11-fp8-flux",
-    "John6666/wai-ani-flux-v10forfp8-fp8-flux",
-    "John6666/flux-dev8-anime-nsfw-fp8-flux",
-    "John6666/nepotism-fuxdevschnell-v3aio-fp8-flux",
-    "John6666/sumeshi-flux1s-v002e-fp8-flux",
-    "John6666/fca-style-v33-x10-8step-fp8-flux",
-    "John6666/lyh-anime-v10f1-fp8-flux",
-    "John6666/lyh-dalle-anime-v12dalle-fp8-flux",
-    "John6666/lyh-anime-flux-v2a1-fp8-flux",
-    "John6666/glimmerkin-flux-cute-v10-fp8-flux",
-    "John6666/niji-style-flux-devfp8-fp8-flux",
-    "John6666/niji56-style-v3-fp8-flux",
-    "John6666/xe-anime-flux-v04-fp8-flux",
-    "John6666/xe-figure-flux-01-fp8-flux",
-    "John6666/xe-pixel-flux-01-fp8-flux",
-    "John6666/xe-guoman-flux-02-fp8-flux",
-    "John6666/carnival-unchained-v10-fp8-flux",
-    "John6666/fluxunchained-artfulnsfw-fut516xfp8e4m3fnv11-fp8-flux",
-    "John6666/fastflux-unchained-t5f16-fp8-flux",
-    "John6666/iniverse-mix-xl-sfwnsfw-fluxdfp16nsfwv11-fp8-flux",
-    "John6666/nsfw-master-flux-lora-merged-with-flux1-dev-fp16-v10-fp8-flux",
-    "John6666/the-araminta-flux1a1-fp8-flux",
-    "John6666/acorn-is-spinning-flux-v11-fp8-flux",
-    "John6666/real-horny-pro-fp8-flux",
-    "John6666/centerfold-flux-v20fp8e5m2-fp8-flux",
-    "John6666/jib-mix-flux-v208stephyper-fp8-flux",
-    "John6666/sapianf-nude-men-women-for-flux-v20fp16-fp8-flux",
-    "John6666/flux-asian-realistic-v10-fp8-flux",
-    "John6666/fluxasiandoll-v10-fp8-flux",
-    "John6666/xe-asian-flux-01-fp8-flux",
-    "John6666/fluxescore-dev-v10fp16-fp8-flux",
-    # "",
-]
-
-
-num_loras = 3
-num_cns = 2
 control_images = [None] * num_cns
 control_modes = [-1] * num_cns
 control_scales = [0] * num_cns
@@ -309,6 +262,12 @@ def get_trigger_word(lorajson: list[dict]):
     return trigger
 
 
+def get_model_trigger(model_name: str):
+    trigger = ""
+    if model_name in model_trigger.keys(): trigger += ", " + model_trigger[model_name]
+    return trigger
+
+
 # https://huggingface.co/docs/diffusers/v0.23.1/en/api/loaders#diffusers.loaders.LoraLoaderMixin.fuse_lora
 # https://github.com/huggingface/diffusers/issues/4919
 def fuse_loras(pipe, lorajson: list[dict]):
@@ -374,6 +333,27 @@ def enhance_prompt(input_prompt):
     return enhanced_text
 
 
+def save_image(image, savefile, modelname, prompt, height, width, steps, cfg, seed):
+    import uuid
+    from PIL import Image, PngImagePlugin
+    import json
+    try:
+        if savefile is None: savefile = f"{modelname.split('/')[-1]}_{str(uuid.uuid4())}.png"
+        metadata = {"prompt": prompt, "Model": {"Model": modelname.split("/")[-1]}}
+        metadata["num_inference_steps"] = steps
+        metadata["guidance_scale"] = cfg
+        metadata["seed"] = seed
+        metadata["resolution"] = f"{width} x {height}"
+        metadata_str = json.dumps(metadata)
+        info = PngImagePlugin.PngInfo()
+        info.add_text("metadata", metadata_str)
+        image.save(savefile, "PNG", pnginfo=info)
+        return str(Path(savefile).resolve())
+    except Exception as e:
+        print(f"Failed to save image file: {e}")
+        raise Exception(f"Failed to save image file:") from e
+
+
 load_prompt_enhancer.zerogpu = True
 fuse_loras.zerogpu = True
 preprocess_image.zerogpu = True

requirements.txt

@@ -15,4 +15,5 @@ opencv-python
 deepspeed
 mediapipe
 openai==1.37.0
-translatepy
+translatepy
+xformers