Upload 31 files

README.md

@@ -5,13 +5,13 @@ colorFrom: purple
 colorTo: red
 sdk: gradio
 sdk_version: 5.45.0
+python_version: "3.10"
 app_file: app.py
 license: mit
 short_description: Text-to-Image
 pinned: true
 preload_from_hub:
   - madebyollin/sdxl-vae-fp16-fix config.json,diffusion_pytorch_model.safetensors
-hf_oauth: true
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -16,7 +16,9 @@ from dc import (infer, _infer, pass_result, get_diffusers_model_list, get_sample
 # Translator
 from llmdolphin import (dolphin_respond_auto, dolphin_parse_simple,
     get_llm_formats, get_dolphin_model_format, get_dolphin_models, get_dolphin_loras, select_dolphin_lora, add_dolphin_loras,
-    get_dolphin_model_info, select_dolphin_model, select_dolphin_format, get_dolphin_sysprompt)
+    get_dolphin_model_info, select_dolphin_model, select_dolphin_format, get_dolphin_sysprompt, initialize_llmdolphin_runtime)
+initialize_llmdolphin_runtime()
+
 # Tagger
 from tagger.v2 import v2_upsampling_prompt, V2_ALL_MODELS
 from tagger.utils import (gradio_copy_text, gradio_copy_prompt, COPY_ACTION_JS,
@@ -681,7 +683,6 @@ with gr.Blocks(fill_width=True, elem_id="container", css=css) as demo:
     with gr.Tab("Preprocessor", render=True):
         preprocessor_tab()
 
-    gr.LoginButton()
     gr.DuplicateButton(value="Duplicate Space for private use (This demo does not work on CPU. Requires GPU Space)")
 
 if __name__ == "__main__":

constants.py

@@ -19,7 +19,7 @@ DOWNLOAD_MODEL = "https://huggingface.co/zuv0/test/resolve/main/milkyWonderland_
 DOWNLOAD_VAE = "https://huggingface.co/Anzhc/Anzhcs-VAEs/resolve/main/SDXL%20Anime%20VAE%20Dec-only%20B3.safetensors, https://huggingface.co/fp16-guy/anything_kl-f8-anime2_vae-ft-mse-840000-ema-pruned_blessed_clearvae_fp16_cleaned/resolve/main/vae-ft-mse-840000-ema-pruned_fp16.safetensors?download=true"
 
 # - **Download LoRAs**
-DOWNLOAD_LORA = "https://huggingface.co/Leopain/color/resolve/main/Coloring_book_-_LineArt.safetensors, https://civitai.com/api/download/models/135867, https://huggingface.co/Linaqruf/anime-detailer-xl-lora/resolve/main/anime-detailer-xl.safetensors?download=true, https://huggingface.co/Linaqruf/style-enhancer-xl-lora/resolve/main/style-enhancer-xl.safetensors?download=true, https://huggingface.co/ByteDance/Hyper-SD/resolve/main/Hyper-SD15-8steps-CFG-lora.safetensors?download=true, https://huggingface.co/ByteDance/Hyper-SD/resolve/main/Hyper-SDXL-8steps-CFG-lora.safetensors?download=true"
+DOWNLOAD_LORA = "https://huggingface.co/Leopain/color/resolve/main/Coloring_book_-_LineArt.safetensors, https://civitai.com/api/download/models/135867, https://huggingface.co/Linaqruf/anime-detailer-xl-lora/resolve/main/anime-detailer-xl.safetensors?download=true, https://huggingface.co/Linaqruf/style-enhancer-xl-lora/resolve/main/style-enhancer-xl.safetensors?download=true"
 
 LOAD_DIFFUSERS_FORMAT_MODEL = [
     'TestOrganizationPleaseIgnore/potato_quality_anime_plzwork_sdxl',

dc.py

@@ -181,6 +181,11 @@ class GuiSD:
         # Avoid duplicate downloads
         self.active_downloads = set()
         self.download_lock = threading.Lock()
+
+        # Anti-abuse: track new model requests.
+        self.used_models = []
+        self.new_model_history = []
+
     def update_storage_models(self, storage_floor_gb=24, required_inventory_for_purge=3):
         while get_used_storage_gb() > storage_floor_gb:
             if len(self.inventory) < required_inventory_for_purge:
@@ -205,6 +210,34 @@ class GuiSD:
         print(self.inventory)
 
     def load_new_model(self, model_name, vae_model, task, controlnet_model, progress=gr.Progress(track_tqdm=True)):
+
+        if model_name != model_list[0]:
+            # --- Anti-Abuse Check Start ---
+            if model_name in self.used_models:
+                # Move to the end to mark as the most recently used.
+                self.used_models.remove(model_name)
+                self.used_models.append(model_name)
+            else:
+                current_time = datetime.now()
+                # Retain history of new model requests from the last 20 minutes.
+                self.new_model_history = [
+                    t for t in self.new_model_history
+                    if (current_time - t).total_seconds() < 1200
+                ]
+
+                # Allow a maximum of 5 new model requests per 20 minutes.
+                if len(self.new_model_history) >= 5:
+                    yield "Rate limit exceeded: Too many new models requested."
+                    raise gr.Error("Too many new models requested. Please reuse your previously loaded models or wait a few minutes before trying new ones.")
+
+                self.new_model_history.append(current_time)
+                self.used_models.append(model_name)
+
+                # Cap the reuse list to the 5 most recent models.
+                if len(self.used_models) > 5:
+                    self.used_models.pop(0)
+            # --- Anti-Abuse Check End ---
+
         lock_key = model_name
 
         while True:

env.py

@@ -77,6 +77,10 @@ LOAD_DIFFUSERS_FORMAT_MODEL = [
     'Raelina/Raehoshi-illust-XL-7',
     'Raelina/Raehoshi-illust-XL-7.1',
     'Raelina/Raehoshi-illust-XL-8',
+    'Raelina/Raehoshi-illust-XL-8.1',
+    'Raelina/Raehoshi-illust-XL-9',
+    'Raelina/Raehoshi-illust-XL-9.1',
+    'Raelina/Raehoshi-illust-vpred',
     'camenduru/FLUX.1-dev-diffusers',
     'black-forest-labs/FLUX.1-schnell',
     'sayakpaul/FLUX.1-merged',

formatter.py

@@ -1,43 +1,71 @@
-from llama_cpp_agent.messages_formatter import MessagesFormatter, PromptMarkers, Roles
-
-mistral_v1_markers = {
-    Roles.system: PromptMarkers(""" [INST]""", """ [/INST] Understood.</s>"""),
-    Roles.user: PromptMarkers(""" [INST]""", """ [/INST]"""),
-    Roles.assistant: PromptMarkers(" ", "</s>"),
-    Roles.tool: PromptMarkers("", ""),
-}
-
-mistral_v1_formatter = MessagesFormatter(
-    pre_prompt="",
-    prompt_markers=mistral_v1_markers,
-    include_sys_prompt_in_first_user_message=False,
-    default_stop_sequences=["</s>"]
-)
-
-mistral_v2_markers = {
-    Roles.system: PromptMarkers("""[INST] """, """[/INST] Understood.</s>"""),
-    Roles.user: PromptMarkers("""[INST] """, """[/INST]"""),
-    Roles.assistant: PromptMarkers(" ", "</s>"),
-    Roles.tool: PromptMarkers("", ""),
-}
-
-mistral_v2_formatter = MessagesFormatter(
-    pre_prompt="",
-    prompt_markers=mistral_v2_markers,
-    include_sys_prompt_in_first_user_message=False,
-    default_stop_sequences=["</s>"]
-)
-
-mistral_v3_tekken_markers = {
-    Roles.system: PromptMarkers("""[INST]""", """[/INST]Understood.</s>"""),
-    Roles.user: PromptMarkers("""[INST]""", """[/INST]"""),
-    Roles.assistant: PromptMarkers("", "</s>"),
-    Roles.tool: PromptMarkers("", ""),
-}
-
-mistral_v3_tekken_formatter = MessagesFormatter(
-    pre_prompt="",
-    prompt_markers=mistral_v3_tekken_markers,
-    include_sys_prompt_in_first_user_message=False,
-    default_stop_sequences=["</s>"]
-)
+class LazyMessagesFormatterType:
+    def __getattr__(self, name):
+        return LazyLlamaAgentValue("MessagesFormatterType", name)
+
+
+class LazyLlamaAgentValue:
+    def __init__(self, source, name):
+        self.source = source
+        self.name = name
+
+    def resolve(self):
+        if self.source == "MessagesFormatterType":
+            from llama_cpp_agent import MessagesFormatterType
+            return getattr(MessagesFormatterType, self.name)
+        if self.source == "custom_formatter":
+            return build_mistral_formatter(self.name)
+        raise ValueError(f"Unknown lazy llama agent value source: {self.source}")
+
+    def __repr__(self):
+        return f"<LazyLlamaAgentValue {self.source}.{self.name}>"
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, LazyLlamaAgentValue)
+            and self.source == other.source
+            and self.name == other.name
+        )
+
+    def __hash__(self):
+        return hash((self.source, self.name))
+
+
+MessagesFormatterType = LazyMessagesFormatterType()
+
+
+def build_mistral_formatter(name):
+    from llama_cpp_agent.messages_formatter import MessagesFormatter, PromptMarkers, Roles
+
+    if name == "mistral_v1":
+        markers = {
+            Roles.system: PromptMarkers(""" [INST]""", """ [/INST]"""),
+            Roles.user: PromptMarkers(""" [INST]""", """ [/INST]"""),
+            Roles.assistant: PromptMarkers(""" """, """</s>"""),
+            Roles.tool: PromptMarkers("", ""),
+        }
+        return MessagesFormatter("", markers, False, ["</s>"])
+
+    if name == "mistral_v2":
+        markers = {
+            Roles.system: PromptMarkers("""[INST] """, """[/INST]"""),
+            Roles.user: PromptMarkers("""[INST] """, """[/INST]"""),
+            Roles.assistant: PromptMarkers(""" """, """</s>"""),
+            Roles.tool: PromptMarkers("", ""),
+        }
+        return MessagesFormatter("", markers, False, ["</s>"])
+
+    if name == "mistral_v3_tekken":
+        markers = {
+            Roles.system: PromptMarkers("""[INST]""", """[/INST]"""),
+            Roles.user: PromptMarkers("""[INST]""", """[/INST]"""),
+            Roles.assistant: PromptMarkers("""""", """</s>"""),
+            Roles.tool: PromptMarkers("", ""),
+        }
+        return MessagesFormatter("", markers, False, ["</s>"])
+
+    raise ValueError(f"Unknown Mistral formatter: {name}")
+
+
+mistral_v1_formatter = LazyLlamaAgentValue("custom_formatter", "mistral_v1")
+mistral_v2_formatter = LazyLlamaAgentValue("custom_formatter", "mistral_v2")
+mistral_v3_tekken_formatter = LazyLlamaAgentValue("custom_formatter", "mistral_v3_tekken")

llmenv.py

@@ -1,8 +1,11 @@
-from llama_cpp_agent import MessagesFormatterType
-from formatter import mistral_v1_formatter, mistral_v2_formatter, mistral_v3_tekken_formatter
+from formatter import MessagesFormatterType, mistral_v1_formatter, mistral_v2_formatter, mistral_v3_tekken_formatter
 from pathlib import Path
 
 
+GRADIO_DEBUG_ENV_NAME = "GRADIO_DEBUG"
+LLMDOLPHIN_STATE_NAMESPACE = "__llmdolphin__"
+
+
 llm_models = {
     #"": ["", MessagesFormatterType.LLAMA_3],
     #"": ["", MessagesFormatterType.MISTRAL],
@@ -161,8 +164,30 @@ llm_models = {
     "claude-3.7-sonnet-reasoning-gemma3-12B.Q4_K_M.gguf": ["mradermacher/claude-3.7-sonnet-reasoning-gemma3-12B-GGUF", MessagesFormatterType.ALPACA],
     "allura-org_MN-Lyrebird-12B-Q4_K_M.gguf": ["bartowski/allura-org_MN-Lyrebird-12B-GGUF", MessagesFormatterType.MISTRAL],
     "ape-fiction-2-mistral-nemo.Q4_K_M.gguf": ["mradermacher/ape-fiction-2-mistral-nemo-GGUF", MessagesFormatterType.MISTRAL],
+    "Shadow-Crystal-12B.Q4_K_M.gguf": ["mradermacher/Shadow-Crystal-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "Starry-Shadow-12B.Q4_K_M.gguf": ["mradermacher/Starry-Shadow-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "Radiant-Shadow-12B.Q4_K_M.gguf": ["mradermacher/Radiant-Shadow-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "Arsenic-Shahrazad-12B-v3.i1-Q4_K_M.gguf": ["mradermacher/Arsenic-Shahrazad-12B-v3-i1-GGUF", MessagesFormatterType.MISTRAL],
+    "Mistral-Nemo-2407-Instruct-12B-Deep-Thinking-Claude-Gemini-GPT5.2.i1-Q4_K_M.gguf": ["mradermacher/Mistral-Nemo-2407-Instruct-12B-Deep-Thinking-Claude-Gemini-GPT5.2-i1-GGUF", MessagesFormatterType.MISTRAL],
+    "Lunar-Abyss-12B.Q4_K_S.gguf": ["mradermacher/Lunar-Abyss-12B-GGUF", MessagesFormatterType.CHATML],
+    "prototype-x-12b-q6_k.gguf": ["Vortex5/Prototype-X-12b-Q6_K-GGUF", MessagesFormatterType.MISTRAL],
+    "gemma-3n-E4B-it-PaperWitch-heresy.i1-Q4_K_M.gguf": ["mradermacher/gemma-3n-E4B-it-PaperWitch-heresy-i1-GGUF", MessagesFormatterType.ALPACA],
+    "Violet-Mist-12B.i1-Q4_K_M.gguf": ["mradermacher/Violet-Mist-12B-i1-GGUF", MessagesFormatterType.MISTRAL],
+    "TheStoryteller-12B.i1-Q4_K_M.gguf": ["mradermacher/TheStoryteller-12B-i1-GGUF", MessagesFormatterType.CHATML],
+    "Dreamstar-12B.Q4_K_M.gguf": ["mradermacher/Dreamstar-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "Mistral-Nemo-Instruct-2407-heretic-noslop-MPOA.Q4_K_M.gguf": ["mradermacher/Mistral-Nemo-Instruct-2407-heretic-noslop-MPOA-GGUF", MessagesFormatterType.MISTRAL],
+    "Equatorium-v3-12B.Q4_K_M.gguf": ["mradermacher/Equatorium-v3-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "RPBizkit-v6-12B.Q4_K_M.gguf": ["mradermacher/RPBizkit-v6-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "MN-VelvetCafe-RP-12B.i1-Q4_K_M.gguf": ["mradermacher/MN-VelvetCafe-RP-12B-i1-GGUF", MessagesFormatterType.CHATML],
+    "Aurora-Mirage-12B.Q4_K_M.gguf": ["mradermacher/Aurora-Mirage-12B-GGUF", MessagesFormatterType.CHATML],
+    "littlemonster-reasoning-12B-QKVO-heretic-HF.Q4_K_M.gguf": ["mradermacher/littlemonster-reasoning-12B-QKVO-heretic-HF-GGUF", MessagesFormatterType.ALPACA],
+    "QuasiStarSynth-12B-absolute-heresy.Q4_K_M.gguf": ["mradermacher/QuasiStarSynth-12B-absolute-heresy-GGUF", MessagesFormatterType.CHATML],
+    "ColdBrew-Nemo-12B-Arcane-Fusion-Combined-Thinker.Q4_K_M.gguf": ["mradermacher/ColdBrew-Nemo-12B-Arcane-Fusion-Combined-Thinker-GGUF", MessagesFormatterType.MISTRAL],
+    "AuroSlayerEtherealKrixUnslopMellRPMaxDARKNESS-12B.i1-Q4_K_M.gguf": ["mradermacher/AuroSlayerEtherealKrixUnslopMellRPMaxDARKNESS-12B-i1-GGUF", MessagesFormatterType.MISTRAL],
     "Irixxed_Homunculus-12B-Q3T-v.0.3-Reasoner.Q4_K_M.gguf": ["mradermacher/Irixxed_Homunculus-12B-Q3T-v.0.3-Reasoner-GGUF", MessagesFormatterType.MISTRAL],
     "Gemma-2-Llama-Swallow-9b-pt-v0.1.Q4_K_M.gguf": ["mradermacher/Gemma-2-Llama-Swallow-9b-pt-v0.1-GGUF", MessagesFormatterType.ALPACA],
+    "Tlacuilo-12B.Q4_K_M.gguf": ["mradermacher/Tlacuilo-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "Nemo-Instruct-2407-MPOA-v4-12B.i1-Q4_K_M.gguf": ["mradermacher/Nemo-Instruct-2407-MPOA-v4-12B-i1-GGUF", MessagesFormatterType.MISTRAL],
     "Qwen2.5-7B-base-french-bespoke-stratos-full-sft.i1-Q5_K_S.gguf": ["mradermacher/Qwen2.5-7B-base-french-bespoke-stratos-full-sft-i1-GGUF", MessagesFormatterType.OPEN_CHAT],
     "Protestant-Christian-Bible-Expert-v2.0-12B.Q4_K_M.gguf": ["mradermacher/Protestant-Christian-Bible-Expert-v2.0-12B-GGUF", MessagesFormatterType.MISTRAL],
     "openbuddy-qwen2.5llamaify-14b-v23.1-200k.i1-Q4_K_M.gguf": ["mradermacher/openbuddy-qwen2.5llamaify-14b-v23.1-200k-i1-GGUF", MessagesFormatterType.OPEN_CHAT],
@@ -178,6 +203,29 @@ llm_models = {
     #"": ["", MessagesFormatterType.OPEN_CHAT],
     #"": ["", MessagesFormatterType.CHATML],
     #"": ["", MessagesFormatterType.PHI_3],
+    "Equatorium-v2-12B.Q4_K_M.gguf": ["mradermacher/Equatorium-v2-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "Mangione-12B-Model_Stock.Q4_K_M.gguf": ["mradermacher/Mangione-12B-Model_Stock-GGUF", MessagesFormatterType.MISTRAL],
+    "gemma-3-12b-it-biprojected-abliterated.Q4_K_M.gguf": ["mradermacher/gemma-3-12b-it-biprojected-abliterated-GGUF", MessagesFormatterType.ALPACA],
+    "mistralai-Mistral-Nemo-Instruct-2407-extensive-BP-abliteration-12B.Q4_K_M.gguf": ["mradermacher/mistralai-Mistral-Nemo-Instruct-2407-extensive-BP-abliteration-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "PokeeAI_pokee_research_7b-Q5_K_M.gguf": ["bartowski/PokeeAI_pokee_research_7b-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "FARE-8B.Q5_K_M.gguf": ["mradermacher/FARE-8B-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "AdaVaR-3B.Q5_K_M.gguf": ["mradermacher/AdaVaR-3B-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "Qwen2.5-14B-Valor.Q4_K_M.gguf": ["mradermacher/Qwen2.5-14B-Valor-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "ukko-thinking-3.09b-v1.i1-Q5_K_M.gguf": ["mradermacher/ukko-thinking-3.09b-v1-i1-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "Logos-3B.Q5_K_M.gguf": ["mradermacher/Logos-3B-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "ReSeek-qwen2.5-7b-em-grpo.Q5_K_M.gguf": ["mradermacher/ReSeek-qwen2.5-7b-em-grpo-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "VideoChat-R1_5.Q5_K_M.gguf": ["mradermacher/VideoChat-R1_5-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "zen-eco-4b-instruct.i1-Q5_K_M.gguf": ["mradermacher/zen-eco-4b-instruct-i1-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "nytheria-3b.i1-Q5_K_M.gguf": ["mradermacher/nytheria-3b-i1-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "VideoJudgeR-3B.Q5_K_M.gguf": ["mradermacher/VideoJudgeR-3B-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "Magrathic-12B.Q4_K_M.gguf": ["mradermacher/Magrathic-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "phoenix-core-v1.0.Q4_K_M.gguf": ["mradermacher/phoenix-core-v1.0-GGUF", MessagesFormatterType.MISTRAL],
+    "GUI-G1-3B-v1.Q5_K_M.gguf": ["mradermacher/GUI-G1-3B-v1-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "Impish-Irix-Kitsune.Q4_K_M.gguf": ["mradermacher/Impish-Irix-Kitsune-GGUF", MessagesFormatterType.MISTRAL],
+    "L3.3-Nemoblated-8B-V0.1.i1-Q5_K_M.gguf": ["mradermacher/L3.3-Nemoblated-8B-V0.1-i1-GGUF", MessagesFormatterType.LLAMA_3],
+    "Luna-Karcher-12B.Q4_K_M.gguf": ["mradermacher/Luna-Karcher-12B-GGUF", MessagesFormatterType.MISTRAL],
+    "Chimera-DeepSeek-NSFW-8B.Q5_K_M.gguf": ["mradermacher/Chimera-DeepSeek-NSFW-8B-GGUF", MessagesFormatterType.LLAMA_3],
+    "Toolbox-sft-3B.Q5_K_M.gguf": ["mradermacher/Toolbox-sft-3B-GGUF", MessagesFormatterType.OPEN_CHAT],
     "SauerHuatuoSkyworkDeepWatt-o1-Llama-3.1-8B.Q5_K_M.gguf": ["mradermacher/SauerHuatuoSkyworkDeepWatt-o1-Llama-3.1-8B-GGUF", MessagesFormatterType.LLAMA_3],
     "care-japanese-llama3.1-8b.Q5_K_M.gguf": ["mradermacher/care-japanese-llama3.1-8b-GGUF", MessagesFormatterType.LLAMA_3],
     "UltraPatriMerge-12B.Q4_K_M.gguf": ["mradermacher/UltraPatriMerge-12B-GGUF", MessagesFormatterType.MISTRAL],
@@ -3838,8 +3886,11 @@ llm_loras = {str(Path(u).name): u for u in llm_loras_urls}
 llm_models_dir = "./llm_models"
 llm_loras_dir = "./llm_loras"
 
+LLM_FORMAT_AUTO_GGUF_DEFAULT = "__GGUF_DEFAULT__"
+
 
 llm_formats = {
+    "AUTO (GGUF DEFAULT)": LLM_FORMAT_AUTO_GGUF_DEFAULT,
     "MISTRAL": MessagesFormatterType.MISTRAL,
     "CHATML": MessagesFormatterType.CHATML,
     "VICUNA": MessagesFormatterType.VICUNA,

packages.txt

@@ -1 +1,2 @@
-git-lfs aria2 ffmpeg
+git-lfs
+ffmpeg

requirements.txt

@@ -2,19 +2,17 @@ stablepy==0.6.5
 diffusers
 transformers
 accelerate
-torch==2.5.1
+torch==2.8.0
 numpy<2
 gdown
 opencv-python
 huggingface_hub
-hf_xet
-hf_transfer
+spaces
 scikit-build-core
-https://github.com/abetlen/llama-cpp-python/releases/download/v0.3.16-cu124/llama_cpp_python-0.3.16-cp310-cp310-linux_x86_64.whl
+https://github.com/John6666cat/llama-cpp-python/releases/download/v0.3.35-cu128-Basic-linux-20260412/llama_cpp_python-0.3.35+cu128.basic-cp310-cp310-linux_x86_64.whl
 git+https://github.com/John6666cat/llama-cpp-agent
 pybind11>=2.12
 rapidfuzz
-torchvision
 optimum[onnxruntime]
 #dartrs
 git+https://github.com/John6666cat/dartrs
@@ -24,5 +22,6 @@ wrapt-timeout-decorator
 sentencepiece
 unidecode
 matplotlib-inline
-https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.11/flash_attn-2.8.3+cu124torch2.5-cp310-cp310-linux_x86_64.whl
-pydantic==2.10.6
+mediapipe==0.10.5
+einops
+# pydantic==2.10.6

tagger/fl2sd3longcap.py

@@ -1,49 +1,34 @@
-import spaces
-from transformers import AutoProcessor, AutoModelForCausalLM
 import re
-from PIL import Image 
+
+import spaces
 import torch
+from PIL import Image
+
+from tagger.florence2_compat import generate_florence2, load_florence2_bundle, prepare_florence2_inputs
 
-from transformers.utils import is_flash_attn_2_available
-if not is_flash_attn_2_available():
-    import subprocess
-    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 try:
-    fl_model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).to("cpu").eval()
-    fl_processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
+    fl_model, fl_processor, fl_audit = load_florence2_bundle('gokaygokay/Florence-2-SD3-Captioner', device)
 except Exception as e:
     print(e)
     fl_model = fl_processor = None
+    fl_audit = {"load_error": repr(e)}
+
 
 def fl_modify_caption(caption: str) -> str:
-    """
-    Removes specific prefixes from captions if present, otherwise returns the original caption.
-    Args:
-        caption (str): A string containing a caption.
-    Returns:
-        str: The caption with the prefix removed if it was present, or the original caption.
-    """
-    # Define the prefixes to remove
     prefix_substrings = [
         ('captured from ', ''),
-        ('captured at ', '')
+        ('captured at ', ''),
     ]
-    
-    # Create a regex pattern to match any of the prefixes
     pattern = '|'.join([re.escape(opening) for opening, _ in prefix_substrings])
     replacers = {opening.lower(): replacer for opening, replacer in prefix_substrings}
-    
-    # Function to replace matched prefix with its corresponding replacement
+
     def replace_fn(match):
         return replacers[match.group(0).lower()]
-    
-    # Apply the regex to the caption
+
     modified_caption = re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE)
-    
-    # If the caption was modified, return the modified version; otherwise, return the original
     return modified_caption if modified_caption != caption else caption
 
 
@@ -52,32 +37,32 @@ def fl_run_example(image):
     task_prompt = "<DESCRIPTION>"
     prompt = task_prompt + "Describe this image in great detail."
 
-    # Ensure the image is in RGB mode
     if image.mode != "RGB":
         image = image.convert("RGB")
 
-    fl_model.to(device)
-    inputs = fl_processor(text=prompt, images=image, return_tensors="pt").to(device)
-    generated_ids = fl_model.generate(
+    inputs = prepare_florence2_inputs(fl_processor, fl_model, prompt, image, device)
+    generated_ids = generate_florence2(fl_model,
         input_ids=inputs["input_ids"],
+        attention_mask=inputs.get("attention_mask"),
         pixel_values=inputs["pixel_values"],
-        max_new_tokens=1024,
-        num_beams=3
     )
-    fl_model.to("cpu")
     generated_text = fl_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    parsed_answer = fl_processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
+    parsed_answer = fl_processor.post_process_generation(
+        generated_text,
+        task=task_prompt,
+        image_size=(image.width, image.height),
+    )
     return fl_modify_caption(parsed_answer["<DESCRIPTION>"])
 
 
 def predict_tags_fl2_sd3(image: Image.Image, input_tags: str, algo: list[str]):
     def to_list(s):
         return [x.strip() for x in s.split(",") if not s == ""]
-    
+
     def list_uniq(l):
         return sorted(set(l), key=l.index)
-    
-    if not "Use Florence-2-SD3-Long-Captioner" in algo:
+
+    if "Use Florence-2-SD3-Long-Captioner" not in algo:
         return input_tags
     tag_list = list_uniq(to_list(input_tags) + to_list(fl_run_example(image) + ", "))
     tag_list.remove("")

tagger/florence2_compat.py

@@ -0,0 +1,925 @@
+import copy
+import functools
+import gc
+import json
+import os
+import re
+import types
+from typing import Any, Dict, Iterable, Optional, Tuple
+
+import torch
+import transformers
+from transformers import AutoProcessor
+from transformers.tokenization_utils_base import AddedToken
+
+
+try:
+    from huggingface_hub import snapshot_download
+except Exception:
+    snapshot_download = None
+
+try:
+    from safetensors.torch import load_file as load_safetensors_file
+except Exception:
+    load_safetensors_file = None
+
+try:
+    from transformers import AutoModelForCausalLM, Florence2ForConditionalGeneration
+except Exception:
+    AutoModelForCausalLM = None
+    Florence2ForConditionalGeneration = None
+
+try:
+    from transformers.generation.utils import GenerationMixin
+except Exception:
+    GenerationMixin = None
+
+try:
+    from transformers.models.auto.auto_factory import add_generation_mixin_to_remote_model
+except Exception:
+    add_generation_mixin_to_remote_model = None
+
+
+_IMAGE_TOKEN = "<image>"
+
+# Florence2 compat settings. Keep these as plain module constants for easier HF Space maintenance.
+FLORENCE2_LEGACY_ATTN_IMPLEMENTATION = "eager"
+FLORENCE2_LEGACY_USE_CACHE = True
+FLORENCE2_LEGACY_MAX_NUM_BEAMS = 1
+FLORENCE2_PROCESSOR_USE_FAST = False
+FLORENCE2_GENERATE_DEFAULT_MAX_NEW_TOKENS = 512
+FLORENCE2_GENERATE_DEFAULT_NUM_BEAMS = 3
+FLORENCE2_GENERATE_MAX_NEW_TOKENS_BY_REPO = {
+    "MiaoshouAI/Florence-2-large-PromptGen-v2.0": 256,
+    "gokaygokay/Florence-2-SD3-Captioner": 384,
+    "gokaygokay/Florence-2-Flux": 384,
+    "gokaygokay/Florence-2-Flux-Large": 384,
+    "thwri/CogFlorence-2.2-Large": 512,
+}
+
+# Pin the currently working Florence checkpoints so trust_remote_code fetches stay reproducible.
+FLORENCE2_REVISIONS = {
+    "gokaygokay/Florence-2-Flux": "d17454c4e4bccb3bfe52477634089b8515b1bc3c",
+    "gokaygokay/Florence-2-Flux-Large": "ed3af3df6d23d9f25d1dd4ce05ba95bb43c37209",
+    "gokaygokay/Florence-2-SD3-Captioner": "178fd612533079793c33ede7245b29d23db307b5",
+    "thwri/CogFlorence-2.2-Large": "19f2c614fdfd18ba49c81d567d70d3a68313e0bb",
+    "MiaoshouAI/Florence-2-large-PromptGen-v2.0": "4aa33eaf50aab040fe8523312ff52eb53322c220",
+}
+
+# Cog processor is the only unstable case. Try the model repo first as documented,
+# then fall back to upstream Florence processor sources.
+FLORENCE2_PROCESSOR_CANDIDATES = {
+    "thwri/CogFlorence-2.2-Large": [
+        {
+            "repo_id": "thwri/CogFlorence-2.2-Large",
+            "trust_remote_code": True,
+            "revision": "19f2c614fdfd18ba49c81d567d70d3a68313e0bb",
+        },
+        {
+            "repo_id": "microsoft/Florence-2-large",
+            "trust_remote_code": True,
+            "revision": "00d2f1570b00c6dea5df998f5635db96840436bc",
+        },
+        {
+            "repo_id": "microsoft/Florence-2-large",
+            "trust_remote_code": False,
+        },
+    ],
+}
+
+
+def _major_version() -> int:
+    m = re.match(r"(\d+)", getattr(transformers, "__version__", "0"))
+    return int(m.group(1)) if m else 0
+
+
+def _dtype_kwargs(device: str) -> Dict[str, Any]:
+    dtype = torch.float16 if device == "cuda" else torch.float32
+    if _major_version() >= 5:
+        return {"dtype": dtype}
+    return {"torch_dtype": dtype}
+
+
+def _revision_kwargs(repo_id: str) -> Dict[str, Any]:
+    revision = FLORENCE2_REVISIONS.get(repo_id)
+    if not revision:
+        return {}
+    return {"revision": revision, "code_revision": revision}
+
+
+def _processor_candidates(repo_id: str):
+    if repo_id in FLORENCE2_PROCESSOR_CANDIDATES:
+        return FLORENCE2_PROCESSOR_CANDIDATES[repo_id]
+    base = {"repo_id": repo_id, "trust_remote_code": True}
+    base.update(_revision_kwargs(repo_id))
+    return [base, {"repo_id": repo_id, "trust_remote_code": False}]
+
+
+def _processor_use_fast(processor_repo_id: str, requested_repo_id: str) -> bool:
+    if requested_repo_id == "thwri/CogFlorence-2.2-Large":
+        return True
+    return FLORENCE2_PROCESSOR_USE_FAST
+
+
+
+@functools.lru_cache(maxsize=4)
+def _ensure_florence2_runtime(device: str) -> Dict[str, Any]:
+    status: Dict[str, Any] = {
+        "device": device,
+        "flash_attn_2_available": False,
+        "install_attempted": False,
+        "install_succeeded": False,
+        "install_returncode": None,
+    }
+    if device != "cuda":
+        print("[Florence2 compat] runtime probe: non-cuda device")
+    else:
+        print("[Florence2 compat] runtime probe: flash-attn runtime install disabled")
+    return status
+
+def _native_attn_kwargs(device: str) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    runtime = dict(_ensure_florence2_runtime(device))
+    runtime["attn_implementation"] = "default"
+    return {}, runtime
+
+
+def _legacy_attn_kwargs(device: str) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    runtime = dict(_ensure_florence2_runtime(device))
+    requested = str(FLORENCE2_LEGACY_ATTN_IMPLEMENTATION).strip().lower() or "eager"
+    if requested != "eager":
+        print(f"[Florence2 compat] legacy attention request {requested} unsupported here, falling back to eager")
+        requested = "eager"
+    runtime["attn_implementation"] = requested
+    return {"attn_implementation": requested}, runtime
+
+
+def _native_florence_module(model: Any) -> bool:
+    module = getattr(getattr(model, "__class__", None), "__module__", "") or ""
+    return module.startswith("transformers.models.florence2")
+
+
+def _normalize_load_result(result: Any) -> Tuple[Any, Dict[str, Any]]:
+    if isinstance(result, tuple) and len(result) == 2 and isinstance(result[1], dict):
+        return result
+    return result, {}
+
+
+def _safe_to_device(model: Any, device: str):
+    return model.eval().to(device)
+
+
+def _load_processor(repo_id: str):
+    errors = []
+    for candidate in _processor_candidates(repo_id):
+        processor_repo_id = candidate["repo_id"]
+        trust_remote_code = bool(candidate.get("trust_remote_code", True))
+        revision = candidate.get("revision", None)
+        use_fast = _processor_use_fast(processor_repo_id, repo_id)
+        kwargs = {"use_fast": use_fast}
+        if revision:
+            kwargs["revision"] = revision
+        try:
+            processor = AutoProcessor.from_pretrained(
+                processor_repo_id,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            )
+            mode = "trust_remote_code=True" if trust_remote_code else "native fallback"
+            print(
+                f"[Florence2 compat] processor={repo_id} source={processor_repo_id} "
+                f"mode={mode} revision={revision or 'main'} use_fast={use_fast}"
+            )
+            return processor
+        except Exception as e:
+            errors.append({
+                "source": processor_repo_id,
+                "trust_remote_code": trust_remote_code,
+                "revision": revision,
+                "error": repr(e),
+            })
+    raise RuntimeError(f"Failed to load processor for {repo_id}: {errors}")
+
+
+def _load_native_model(repo_id: str, device: str):
+    if Florence2ForConditionalGeneration is None:
+        raise RuntimeError("Florence2ForConditionalGeneration is unavailable in this transformers build")
+    attn_kwargs, runtime = _native_attn_kwargs(device)
+    load_kwargs = dict(output_loading_info=True, **_dtype_kwargs(device), **attn_kwargs, **_revision_kwargs(repo_id))
+    model, loading_info = _normalize_load_result(
+        Florence2ForConditionalGeneration.from_pretrained(repo_id, **load_kwargs)
+    )
+    print(
+        f"[Florence2 compat] model={repo_id} loaded as native Florence2ForConditionalGeneration "
+        f"revision={FLORENCE2_REVISIONS.get(repo_id) or 'main'} attn={runtime.get('attn_implementation')}"
+    )
+    return model, loading_info, runtime
+
+
+def _patch_generation_mixin(module: Any, label: str) -> bool:
+    if module is None or GenerationMixin is None:
+        return False
+    if hasattr(module, "generate") and callable(getattr(module, "generate")):
+        return False
+
+    cls = module.__class__
+    if issubclass(cls, GenerationMixin):
+        return False
+
+    if add_generation_mixin_to_remote_model is not None:
+        try:
+            patched_cls = add_generation_mixin_to_remote_model(cls)
+            if patched_cls is not cls:
+                module.__class__ = patched_cls
+                print(f"[Florence2 compat] patched {label} class with GenerationMixin: {patched_cls.__name__}")
+                return True
+        except Exception as e:
+            print(f"[Florence2 compat] official GenerationMixin patch failed for {label}: {e!r}")
+
+    patched_cls = type(f"{cls.__name__}WithGenerationMixin", (cls, GenerationMixin), {})
+    module.__class__ = patched_cls
+    print(f"[Florence2 compat] patched {label} class with GenerationMixin: {patched_cls.__name__}")
+    return True
+
+
+def _first_non_none(*values):
+    for value in values:
+        if value is not None:
+            return value
+    return None
+
+
+def _token_like_int(value: Any) -> Optional[int]:
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, int):
+        return int(value)
+    return None
+
+
+def _resolve_generation_ids(model: Any, processor: Any = None) -> Dict[str, Optional[int]]:
+    config = getattr(model, "config", None)
+    text_config = getattr(config, "text_config", None) if config is not None else None
+    tokenizer = getattr(processor, "tokenizer", None) if processor is not None else None
+
+    bos_token_id = _first_non_none(
+        _token_like_int(getattr(config, "bos_token_id", None)),
+        _token_like_int(getattr(text_config, "bos_token_id", None)),
+        _token_like_int(getattr(tokenizer, "bos_token_id", None)),
+    )
+    eos_token_id = _first_non_none(
+        _token_like_int(getattr(config, "eos_token_id", None)),
+        _token_like_int(getattr(text_config, "eos_token_id", None)),
+        _token_like_int(getattr(tokenizer, "eos_token_id", None)),
+    )
+    pad_token_id = _first_non_none(
+        _token_like_int(getattr(config, "pad_token_id", None)),
+        _token_like_int(getattr(text_config, "pad_token_id", None)),
+        _token_like_int(getattr(tokenizer, "pad_token_id", None)),
+    )
+    decoder_start_token_id = _first_non_none(
+        _token_like_int(getattr(config, "decoder_start_token_id", None)),
+        _token_like_int(getattr(text_config, "decoder_start_token_id", None)),
+        eos_token_id,
+        bos_token_id,
+    )
+
+    return {
+        "bos_token_id": bos_token_id,
+        "eos_token_id": eos_token_id,
+        "pad_token_id": pad_token_id,
+        "decoder_start_token_id": decoder_start_token_id,
+    }
+
+
+def _apply_generation_ids(target: Any, ids: Dict[str, Optional[int]]) -> None:
+    if target is None:
+        return
+    for key, value in ids.items():
+        if value is None:
+            continue
+        try:
+            current = getattr(target, key, None)
+        except Exception:
+            current = None
+        if current is None:
+            try:
+                setattr(target, key, value)
+            except Exception:
+                pass
+
+
+def _patch_prepare_inputs_for_generation(module: Any) -> bool:
+    if module is None or getattr(module, "_florence2_prepare_inputs_patched", False):
+        return False
+    original = getattr(module, "prepare_inputs_for_generation", None)
+    if not callable(original):
+        return False
+
+    def _wrapped_prepare_inputs_for_generation(self, decoder_input_ids, *args, **kwargs):
+        past_key_values = kwargs.get("past_key_values", None)
+        if past_key_values is not None:
+            reset_to_none = False
+            try:
+                first_layer = past_key_values[0]
+                first_k = first_layer[0] if first_layer is not None else None
+                if first_k is None:
+                    reset_to_none = True
+            except Exception:
+                reset_to_none = True
+            if reset_to_none:
+                kwargs["past_key_values"] = None
+        return original(decoder_input_ids, *args, **kwargs)
+
+    module.prepare_inputs_for_generation = types.MethodType(_wrapped_prepare_inputs_for_generation, module)
+    module._florence2_prepare_inputs_patched = True
+    print("[Florence2 compat] patched legacy prepare_inputs_for_generation guard")
+    return True
+
+
+def _safe_module_device_dtype(module: Any) -> Tuple[Optional[str], Optional[torch.dtype]]:
+    if module is None or not hasattr(module, "parameters"):
+        return None, None
+    try:
+        param = next(module.parameters())
+        return str(param.device), param.dtype
+    except Exception:
+        return None, None
+
+
+
+def _beam_cap(default: int = 1) -> int:
+    raw = str(FLORENCE2_LEGACY_MAX_NUM_BEAMS if FLORENCE2_LEGACY_MAX_NUM_BEAMS is not None else default).strip()
+    try:
+        value = int(raw)
+    except Exception:
+        value = default
+    return max(0, value)
+
+
+def _copy_generation_config(config: Any) -> Any:
+    if config is None:
+        return None
+    try:
+        return copy.deepcopy(config)
+    except Exception:
+        return config
+
+
+def _prepare_generation_config_for_call(model: Any, language_model: Any, kwargs: Dict[str, Any], generation_ids: Dict[str, Optional[int]]) -> None:
+    active_generation_config = _copy_generation_config(getattr(model, "generation_config", None))
+    if active_generation_config is None and language_model is not None:
+        active_generation_config = _copy_generation_config(getattr(language_model, "generation_config", None))
+    if active_generation_config is None:
+        return
+
+    _apply_generation_ids(active_generation_config, generation_ids)
+
+    num_beams = kwargs.get("num_beams", getattr(active_generation_config, "num_beams", None))
+    try:
+        num_beams = int(num_beams) if num_beams is not None else None
+    except Exception:
+        num_beams = None
+
+    if num_beams is not None:
+        active_generation_config.num_beams = num_beams
+
+    if num_beams is None or num_beams <= 1:
+        try:
+            active_generation_config.early_stopping = False
+        except Exception:
+            pass
+        kwargs.pop("early_stopping", None)
+
+    kwargs.setdefault("generation_config", active_generation_config)
+    if language_model is not None:
+        try:
+            language_model.generation_config = active_generation_config
+        except Exception:
+            pass
+
+
+def _patch_legacy_generate_bridge(model: Any, processor: Any = None) -> None:
+    language_model = getattr(model, "language_model", None)
+    patched_inner = _patch_generation_mixin(language_model, "language_model")
+    _patch_prepare_inputs_for_generation(language_model)
+
+    generation_ids = _resolve_generation_ids(model, processor)
+    _apply_generation_ids(getattr(model, "generation_config", None), generation_ids)
+
+    if language_model is not None:
+        try:
+            if hasattr(model, "generation_config"):
+                language_model.generation_config = model.generation_config
+            if hasattr(language_model, "generation_config"):
+                _apply_generation_ids(language_model.generation_config, generation_ids)
+            print("[Florence2 compat] bridged generation_config -> language_model.generation_config")
+        except Exception as e:
+            print(f"[Florence2 compat] failed to bridge generation_config: {e!r}")
+
+    if getattr(model, "_florence2_generate_bridge_patched", False):
+        return
+
+    original_generate = getattr(model, "generate", None)
+    if not callable(original_generate):
+        return
+
+    def _wrapped_generate(self, *args, **kwargs):
+        lm = getattr(self, "language_model", None)
+        if lm is not None and hasattr(self, "generation_config"):
+            try:
+                lm.generation_config = self.generation_config
+            except Exception:
+                pass
+        generation_ids = _resolve_generation_ids(self, processor)
+        _apply_generation_ids(getattr(self, "generation_config", None), generation_ids)
+        if lm is not None and hasattr(lm, "generation_config"):
+            _apply_generation_ids(lm.generation_config, generation_ids)
+
+        kwargs.setdefault("use_cache", FLORENCE2_LEGACY_USE_CACHE)
+
+        beam_cap = _beam_cap(default=1)
+        if beam_cap > 0:
+            requested_beams = kwargs.get("num_beams", None)
+            if requested_beams is None:
+                kwargs["num_beams"] = beam_cap
+            else:
+                try:
+                    requested_beams_int = int(requested_beams)
+                except Exception:
+                    requested_beams_int = requested_beams
+                if isinstance(requested_beams_int, int) and requested_beams_int > beam_cap:
+                    print(
+                        f"[Florence2 compat] clamped legacy num_beams from {requested_beams_int} to {beam_cap}"
+                    )
+                    kwargs["num_beams"] = beam_cap
+
+        for key, value in generation_ids.items():
+            if value is not None:
+                kwargs.setdefault(key, value)
+
+        _prepare_generation_config_for_call(self, lm, kwargs, generation_ids)
+
+        _condensed_generate_log(self, kwargs)
+        return original_generate(*args, **kwargs)
+
+    model.generate = types.MethodType(_wrapped_generate, model)
+    model._florence2_generate_bridge_patched = True
+    print("[Florence2 compat] legacy generate defaults: " + ", ".join(f"{k}={v}" for k, v in generation_ids.items()))
+    if patched_inner:
+        print("[Florence2 compat] patched legacy Florence generate bridge")
+
+
+def _load_legacy_model(repo_id: str, device: str, processor: Any = None):
+    if AutoModelForCausalLM is None:
+        raise RuntimeError("AutoModelForCausalLM is unavailable in this transformers build")
+    attn_kwargs, runtime = _legacy_attn_kwargs(device)
+    load_kwargs = dict(output_loading_info=True, **_dtype_kwargs(device), **attn_kwargs, **_revision_kwargs(repo_id))
+    try:
+        model, loading_info = _normalize_load_result(
+            AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True, **load_kwargs)
+        )
+    except TypeError:
+        load_kwargs.pop("attn_implementation", None)
+        runtime["attn_implementation"] = "legacy_default"
+        model, loading_info = _normalize_load_result(
+            AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True, **load_kwargs)
+        )
+    _patch_legacy_generate_bridge(model, processor=processor)
+    print(
+        f"[Florence2 compat] model={repo_id} loaded via AutoModelForCausalLM + trust_remote_code=True "
+        f"revision={FLORENCE2_REVISIONS.get(repo_id) or 'main'} attn={runtime.get('attn_implementation')}"
+    )
+    return model, loading_info, runtime
+
+
+def _ensure_image_token(processor: Any, model: Any) -> None:
+    tokenizer = getattr(processor, "tokenizer", None)
+    config = getattr(model, "config", None)
+    if tokenizer is None or config is None:
+        return
+
+    tokenizer.image_token = getattr(tokenizer, "image_token", None) or _IMAGE_TOKEN
+
+    token_id = tokenizer.convert_tokens_to_ids(tokenizer.image_token)
+    needs_add = token_id is None
+    if token_id is not None:
+        unk_id = getattr(tokenizer, "unk_token_id", None)
+        if unk_id is not None and token_id == unk_id:
+            pieces = tokenizer.encode(tokenizer.image_token, add_special_tokens=False)
+            needs_add = len(pieces) != 1 or pieces[0] == unk_id
+
+    added = 0
+    if needs_add:
+        added = tokenizer.add_tokens(
+            AddedToken(tokenizer.image_token, special=True, normalized=False),
+            special_tokens=True,
+        )
+        print(f"[Florence2 compat] added special image token to tokenizer: added={added}")
+
+    pieces = tokenizer.encode(tokenizer.image_token, add_special_tokens=False)
+    if len(pieces) == 1:
+        image_token_id = int(pieces[0])
+        tokenizer.image_token_id = image_token_id
+        if getattr(config, "image_token_id", None) != image_token_id:
+            config.image_token_id = image_token_id
+            print(f"[Florence2 compat] model.config.image_token_id set to {image_token_id}")
+
+    if added > 0 and hasattr(model, "resize_token_embeddings"):
+        try:
+            model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=64)
+        except TypeError:
+            try:
+                model.resize_token_embeddings(len(tokenizer), 64)
+            except TypeError:
+                model.resize_token_embeddings(len(tokenizer))
+        print(f"[Florence2 compat] resize_token_embeddings({len(tokenizer)}) applied")
+
+
+def audit_loading_info(repo_id: str, loading_info: Dict[str, Any]) -> Dict[str, Any]:
+    missing = list(loading_info.get("missing_keys") or [])
+    unexpected = list(loading_info.get("unexpected_keys") or [])
+
+    critical_patterns = (
+        "lm_head.weight",
+        "model.language_model.decoder.",
+        "model.multi_modal_projector.",
+    )
+    legacy_patterns = (
+        "language_model.model.",
+        "language_model.lm_head.weight",
+        "image_projection",
+        "image_proj_norm.",
+        "image_pos_embed.",
+        "visual_temporal_embed.",
+        "vision_tower.",
+    )
+
+    critical_missing = [k for k in missing if any(p in k for p in critical_patterns)]
+    legacy_unexpected = [k for k in unexpected if any(p in k for p in legacy_patterns)]
+    suspicious = bool(critical_missing and legacy_unexpected)
+
+    if suspicious:
+        print(
+            f"[Florence2 compat] suspicious native load mismatch for {repo_id}: "
+            f"critical_missing={len(critical_missing)} legacy_unexpected={len(legacy_unexpected)}"
+        )
+        print("[Florence2 compat] trying legacy->native key conversion fallback")
+
+    return {
+        "missing_keys": missing,
+        "unexpected_keys": unexpected,
+        "missing_count": len(missing),
+        "unexpected_count": len(unexpected),
+        "critical_missing_count": len(critical_missing),
+        "legacy_unexpected_count": len(legacy_unexpected),
+        "suspicious_mismatch": suspicious,
+    }
+
+
+def _collect_weight_files(snapshot_dir: str) -> Iterable[str]:
+    candidates = []
+    for index_name in ("model.safetensors.index.json", "pytorch_model.bin.index.json"):
+        index_path = os.path.join(snapshot_dir, index_name)
+        if os.path.exists(index_path):
+            with open(index_path, "r", encoding="utf-8") as f:
+                index_data = json.load(f)
+            weight_map = index_data.get("weight_map") or {}
+            seen = set()
+            for rel_path in weight_map.values():
+                if rel_path not in seen:
+                    candidates.append(os.path.join(snapshot_dir, rel_path))
+                    seen.add(rel_path)
+            return candidates
+
+    for file_name in (
+        "model.safetensors",
+        "pytorch_model.bin",
+        "pytorch_model.pt",
+    ):
+        path = os.path.join(snapshot_dir, file_name)
+        if os.path.exists(path):
+            return [path]
+
+    for name in sorted(os.listdir(snapshot_dir)):
+        if name.endswith(".safetensors") or name.endswith(".bin"):
+            candidates.append(os.path.join(snapshot_dir, name))
+    return candidates
+
+
+def _load_state_dict_from_snapshot(repo_id: str) -> Dict[str, torch.Tensor]:
+    if snapshot_download is None:
+        raise RuntimeError("huggingface_hub.snapshot_download is unavailable")
+
+    snapshot_dir = snapshot_download(
+        repo_id,
+        revision=FLORENCE2_REVISIONS.get(repo_id),
+        allow_patterns=[
+            "*.safetensors",
+            "*.safetensors.index.json",
+            "*.bin",
+            "*.bin.index.json",
+            "config.json",
+            "generation_config.json",
+        ],
+        ignore_patterns=["*.msgpack", "*.h5", "*.ot", "*.onnx", "*.tflite"],
+    )
+
+    state_dict: Dict[str, torch.Tensor] = {}
+    weight_files = list(_collect_weight_files(snapshot_dir))
+    if not weight_files:
+        raise RuntimeError(f"No model weight files found in snapshot for {repo_id}")
+
+    for path in weight_files:
+        if path.endswith(".safetensors"):
+            if load_safetensors_file is None:
+                raise RuntimeError("safetensors is unavailable but checkpoint uses safetensors")
+            shard = load_safetensors_file(path, device="cpu")
+        else:
+            shard = torch.load(path, map_location="cpu")
+            if isinstance(shard, dict) and "state_dict" in shard and isinstance(shard["state_dict"], dict):
+                shard = shard["state_dict"]
+        state_dict.update(shard)
+
+    print(f"[Florence2 compat] loaded raw state_dict from snapshot: files={len(weight_files)} tensors={len(state_dict)}")
+    return state_dict
+
+
+def _rename_vision_tower_key(key: str) -> str:
+    new_key = f"model.{key}"
+    new_key = new_key.replace(".convs.", ".convs.")
+    new_key = new_key.replace(".proj.", ".conv.") if ".convs." in new_key else new_key
+
+    replacements = (
+        (".spatial_block.conv1.fn.dw.", ".spatial_block.conv1."),
+        (".spatial_block.window_attn.norm.", ".spatial_block.norm1."),
+        (".spatial_block.window_attn.fn.", ".spatial_block.window_attn."),
+        (".spatial_block.ffn.norm.", ".spatial_block.norm2."),
+        (".spatial_block.ffn.fn.net.", ".spatial_block.ffn."),
+        (".channel_block.conv1.fn.dw.", ".channel_block.conv1."),
+        (".channel_block.channel_attn.norm.", ".channel_block.norm1."),
+        (".channel_block.channel_attn.fn.", ".channel_block.channel_attn."),
+        (".channel_block.conv2.fn.dw.", ".channel_block.conv2."),
+        (".channel_block.ffn.norm.", ".channel_block.norm2."),
+        (".channel_block.ffn.fn.net.", ".channel_block.ffn."),
+    )
+    for src, dst in replacements:
+        if src in new_key:
+            new_key = new_key.replace(src, dst)
+    return new_key
+
+
+def convert_legacy_state_dict_to_native(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    converted: Dict[str, torch.Tensor] = {}
+
+    direct_map = {
+        "image_proj_norm.weight": "model.multi_modal_projector.image_proj_norm.weight",
+        "image_proj_norm.bias": "model.multi_modal_projector.image_proj_norm.bias",
+        "image_pos_embed.row_embeddings.weight": "model.multi_modal_projector.image_position_embed.row_embeddings.weight",
+        "image_pos_embed.column_embeddings.weight": "model.multi_modal_projector.image_position_embed.column_embeddings.weight",
+        "visual_temporal_embed.pos_idx_to_embed": "model.multi_modal_projector.visual_temporal_embed.pos_idx_to_embed",
+        "language_model.lm_head.weight": "lm_head.weight",
+    }
+
+    for key, value in state_dict.items():
+        tensor = value
+        new_key = key
+
+        if key in direct_map:
+            new_key = direct_map[key]
+        elif key == "image_projection":
+            new_key = "model.multi_modal_projector.image_projection.weight"
+            tensor = value.transpose(1, 0).contiguous()
+        elif key.startswith("language_model.model."):
+            new_key = key.replace("language_model.model.", "model.language_model.", 1)
+        elif key.startswith("vision_tower."):
+            new_key = _rename_vision_tower_key(key)
+        elif key.startswith("model.") or key == "lm_head.weight":
+            new_key = key
+        else:
+            # Keep unknown keys as-is; they will appear in audit logs if unused.
+            new_key = key
+
+        converted[new_key] = tensor
+
+    print(
+        "[Florence2 compat] converted legacy state_dict to native-style keys: "
+        f"input={len(state_dict)} output={len(converted)}"
+    )
+    return converted
+
+
+def _rebuild_native_model_from_converted_state(base_model: Any, processor: Any, repo_id: str):
+    state_dict = _load_state_dict_from_snapshot(repo_id)
+    converted_state = convert_legacy_state_dict_to_native(state_dict)
+
+    native_cls = Florence2ForConditionalGeneration or type(base_model)
+    converted_model = native_cls(base_model.config)
+    incompatible = converted_model.load_state_dict(converted_state, strict=False)
+    if hasattr(converted_model, "tie_weights"):
+        converted_model.tie_weights()
+
+    loading_info = {
+        "missing_keys": list(getattr(incompatible, "missing_keys", []) or []),
+        "unexpected_keys": list(getattr(incompatible, "unexpected_keys", []) or []),
+    }
+    _ensure_image_token(processor, converted_model)
+    audit = audit_loading_info(repo_id, loading_info)
+    audit["converted_from_legacy"] = True
+    return converted_model, audit
+
+
+def _boolish(value: Any) -> Optional[bool]:
+    if isinstance(value, bool):
+        return value
+    return None
+
+
+def _effective_generate_settings(model: Any, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    generation_config = kwargs.get("generation_config", getattr(model, "generation_config", None))
+    use_cache = kwargs.get("use_cache", getattr(generation_config, "use_cache", None))
+    num_beams = kwargs.get("num_beams", getattr(generation_config, "num_beams", None))
+    early_stopping = kwargs.get("early_stopping", getattr(generation_config, "early_stopping", None))
+    max_new_tokens = kwargs.get("max_new_tokens", getattr(generation_config, "max_new_tokens", None))
+    return {
+        "use_cache": use_cache,
+        "num_beams": num_beams,
+        "early_stopping": early_stopping,
+        "max_new_tokens": max_new_tokens,
+    }
+
+
+def _condensed_generate_log(model: Any, kwargs: Dict[str, Any]) -> None:
+    settings = _effective_generate_settings(model, kwargs)
+    loader = getattr(model, "_florence2_loader", "unknown")
+    attn = getattr(model, "_florence2_attn_impl", "unknown")
+    vision_device, vision_dtype = _safe_module_device_dtype(getattr(model, "vision_tower", None))
+    language_device, language_dtype = _safe_module_device_dtype(getattr(model, "language_model", None))
+    print(
+        "[Florence2 compat] call "
+        f"loader={loader} attn={attn} "
+        f"vt={vision_device}/{vision_dtype} lm={language_device}/{language_dtype} "
+        f"cache={settings['use_cache']} beams={settings['num_beams']} "
+        f"early={settings['early_stopping']} max_new_tokens={settings['max_new_tokens']} "
+        f"repo={getattr(model, '_florence2_repo_id', None)}"
+    )
+
+
+def _default_max_new_tokens(model: Any) -> int:
+    repo_id = getattr(model, "_florence2_repo_id", None)
+    if isinstance(repo_id, str) and repo_id in FLORENCE2_GENERATE_MAX_NEW_TOKENS_BY_REPO:
+        return int(FLORENCE2_GENERATE_MAX_NEW_TOKENS_BY_REPO[repo_id])
+    return int(FLORENCE2_GENERATE_DEFAULT_MAX_NEW_TOKENS)
+
+
+def _autocast_context(model: Any):
+    device, dtype = _infer_tensor_runtime(model, "cpu")
+    if not str(device).startswith("cuda"):
+        return torch.autocast(device_type="cpu", enabled=False)
+    target_dtype = dtype if dtype in (torch.float16, torch.bfloat16) else torch.float16
+    return torch.autocast(device_type="cuda", dtype=target_dtype)
+
+
+def generate_florence2(model: Any, **kwargs):
+    kwargs.setdefault("max_new_tokens", _default_max_new_tokens(model))
+    kwargs.setdefault("num_beams", FLORENCE2_GENERATE_DEFAULT_NUM_BEAMS)
+    if isinstance(kwargs.get("num_beams"), int) and kwargs["num_beams"] <= 1:
+        kwargs.pop("early_stopping", None)
+    _condensed_generate_log(model, kwargs)
+    with torch.inference_mode():
+        with _autocast_context(model):
+            return model.generate(**kwargs)
+
+
+def load_florence2_bundle(repo_id: str, device: str, prefer_native_on_v5: bool = True):
+    major = _major_version()
+    processor = _load_processor(repo_id)
+
+    if major < 5 or not prefer_native_on_v5 or Florence2ForConditionalGeneration is None:
+        model, loading_info, runtime = _load_legacy_model(repo_id, device=device, processor=processor)
+        model = _safe_to_device(model, device)
+        audit = audit_loading_info(repo_id, loading_info)
+        audit["loader"] = "legacy_remote_code"
+        audit["runtime"] = runtime
+        model._florence2_loader = audit["loader"]
+        model._florence2_attn_impl = runtime.get("attn_implementation")
+        model._florence2_repo_id = repo_id
+        return model, processor, audit
+
+    model, loading_info, runtime = _load_native_model(repo_id, device=device)
+    audit = audit_loading_info(repo_id, loading_info)
+    audit["loader"] = "native"
+    audit["runtime"] = runtime
+
+    if audit.get("suspicious_mismatch"):
+        try:
+            converted_model, converted_audit = _rebuild_native_model_from_converted_state(model, processor, repo_id)
+            del model
+            gc.collect()
+            model = converted_model
+            audit = converted_audit
+            audit["loader"] = "converted_native"
+            audit["runtime"] = runtime
+        except Exception as e:
+            audit["conversion_error"] = repr(e)
+            print(f"[Florence2 compat] legacy->native conversion fallback failed for {repo_id}: {e!r}")
+
+    if _native_florence_module(model):
+        _ensure_image_token(processor, model)
+
+    model = _safe_to_device(model, device)
+    model._florence2_loader = audit.get("loader", "unknown")
+    model._florence2_attn_impl = runtime.get("attn_implementation")
+    model._florence2_repo_id = repo_id
+    return model, processor, audit
+
+
+def _candidate_image_seq_length(processor: Any, model: Any) -> Optional[int]:
+    candidates = (
+        getattr(getattr(model, "config", None), "image_seq_length", None),
+        getattr(processor, "num_image_tokens", None),
+        getattr(processor, "image_seq_length", None),
+        getattr(getattr(processor, "image_processor", None), "image_seq_length", None),
+    )
+    for value in candidates:
+        if isinstance(value, int) and value > 0:
+            return value
+    return None
+
+
+def _infer_tensor_runtime(model: Any, fallback_device: str) -> Tuple[str, Optional[torch.dtype]]:
+    modules = [
+        getattr(model, "vision_tower", None),
+        getattr(getattr(model, "model", None), "vision_tower", None),
+        model,
+    ]
+    for module in modules:
+        if module is None or not hasattr(module, "parameters"):
+            continue
+        try:
+            param = next(module.parameters())
+            return str(param.device), param.dtype
+        except StopIteration:
+            continue
+        except Exception:
+            continue
+    return fallback_device, getattr(model, "dtype", None)
+
+
+
+def prepare_florence2_inputs(processor: Any, model: Any, prompt: str, image: Any, device: str):
+    inputs = processor(text=prompt, images=image, return_tensors="pt")
+
+    if _native_florence_module(model):
+        input_ids = inputs.get("input_ids")
+        pixel_values = inputs.get("pixel_values")
+        image_token_id = getattr(getattr(model, "config", None), "image_token_id", None)
+        image_seq_length = _candidate_image_seq_length(processor, model)
+
+        if input_ids is not None and pixel_values is not None and image_token_id is not None and image_seq_length:
+            n_image_tokens = int((input_ids == image_token_id).sum().item())
+            if n_image_tokens == 0:
+                batch_size = input_ids.shape[0]
+                prefix = torch.full(
+                    (batch_size, int(image_seq_length)),
+                    int(image_token_id),
+                    dtype=input_ids.dtype,
+                    device=input_ids.device,
+                )
+                inputs["input_ids"] = torch.cat([prefix, input_ids], dim=1)
+
+                attention_mask = inputs.get("attention_mask")
+                if attention_mask is not None:
+                    prefix_mask = torch.ones(
+                        (batch_size, int(image_seq_length)),
+                        dtype=attention_mask.dtype,
+                        device=attention_mask.device,
+                    )
+                    inputs["attention_mask"] = torch.cat([prefix_mask, attention_mask], dim=1)
+
+                print(
+                    f"[Florence2 compat] injected {int(image_seq_length)} image placeholder tokens for native Florence-2"
+                )
+
+    target_device, float_dtype = _infer_tensor_runtime(model, device)
+    prepared = {}
+    for k, v in inputs.items():
+        if not hasattr(v, "to"):
+            prepared[k] = v
+            continue
+        if isinstance(v, torch.Tensor) and torch.is_floating_point(v):
+            prepared[k] = v.to(device=target_device, dtype=float_dtype or v.dtype)
+        else:
+            prepared[k] = v.to(device=target_device)
+
+    pv = prepared.get("pixel_values")
+    if isinstance(pv, torch.Tensor):
+        print(
+            "[Florence2 compat] prepared pixel_values "
+            f"device={pv.device} dtype={pv.dtype} target_float_dtype={float_dtype} native={_native_florence_module(model)}"
+        )
+
+    return prepared
+
+
+def safe_load_florence2_bundle(repo_id: str, device: str, prefer_native_on_v5: bool = True):
+    return load_florence2_bundle(repo_id=repo_id, device=device, prefer_native_on_v5=prefer_native_on_v5)

tagger/tagger.py

@@ -8,13 +8,14 @@ from pathlib import Path
 
 WD_MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
 WD_MODEL_NAME = WD_MODEL_NAMES[0]
+WD_MODEL_REVISION = "b09fdef"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 default_device = device
 
 try:
-    wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True).to(default_device).eval()
-    wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+    wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True, revision=WD_MODEL_REVISION, code_revision=WD_MODEL_REVISION).to(default_device).eval()
+    wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True, revision=WD_MODEL_REVISION, code_revision=WD_MODEL_REVISION)
 except Exception as e:
     print(e)
     wd_model = wd_processor = None

utils.py

@@ -275,11 +275,11 @@ def civ_redirect_down(url, dir_, civitai_api_key, romanize, alternative_name):
     elif os.path.exists(os.path.join(dir_, filename_base)):
         return os.path.join(dir_, filename_base), filename_base
 
-    aria2_command = (
-        f'aria2c --console-log-level=error --summary-interval=10 -c -x 16 '
-        f'-k 1M -s 16 -d "{dir_}" -o "{filename_base}" "{redirect_url}"'
+    wget_command = (
+        f'wget -c -nv '
+        f'-O "{os.path.join(dir_, filename_base)}" "{redirect_url}"'
     )
-    r_code = os.system(aria2_command)  # noqa
+    r_code = os.system(wget_command)  # noqa
 
     # if r_code != 0:
     #     raise RuntimeError(f"Failed to download file: {filename_base}. Error code: {r_code}")
@@ -293,27 +293,32 @@ def civ_redirect_down(url, dir_, civitai_api_key, romanize, alternative_name):
 
 def civ_api_down(url, dir_, civitai_api_key, civ_filename):
     """
-    This method is susceptible to being blocked because it generates a lot of temp redirect links with aria2c.
-    If an API key limit is reached, generating a new API key and using it can fix the issue.
+    This method is susceptible to being blocked because it generates a lot of temp redirect links with wget.
+    If an API key limit is reached, generating a new API key and using it can fix the issue. no
     """
     output_path = None
-
+    
     url_dl = url + f"?token={civitai_api_key}"
+    
     if not civ_filename:
-        aria2_command = f'aria2c -c -x 1 -s 1 -d "{dir_}" "{url_dl}"'
-        os.system(aria2_command)
+        wget_command = (
+            f'wget -c -nv '
+            f'-P "{dir_}" "{url_dl}"'
+        )
+        os.system(wget_command)
+    
     else:
         output_path = os.path.join(dir_, civ_filename)
+    
         if not os.path.exists(output_path):
-            aria2_command = (
-                f'aria2c --console-log-level=error --summary-interval=10 -c -x 16 '
-                f'-k 1M -s 16 -d "{dir_}" -o "{civ_filename}" "{url_dl}"'
+            wget_command = (
+                f'wget -c -nv '
+                f'-O "{output_path}" "{url_dl}"'
             )
-            os.system(aria2_command)
-
+            os.system(wget_command)
+    
     return output_path
 
-
 def drive_down(url, dir_):
     import gdown
 
@@ -354,10 +359,16 @@ def hf_down(url, dir_, hf_token, romanize):
         url = url.replace("/blob/", "/resolve/")
 
     if hf_token:
-        user_header = f'"Authorization: Bearer {hf_token}"'
-        os.system(f"aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 {url} -d {dir_}  -o {filename}")
+        os.system(
+            f'wget -c -nv '
+            f'--header="Authorization: Bearer {hf_token}" '
+            f'-O "{os.path.join(dir_, filename)}" "{url}"'
+        )
     else:
-        os.system(f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {dir_}  -o {filename}")
+        os.system(
+            f'wget -c -nv '
+            f'-O "{os.path.join(dir_, filename)}" "{url}"'
+        )
 
     return output_path
 
@@ -370,7 +381,8 @@ def download_things(directory, url, hf_token="", civitai_api_key="", romanize=Fa
         downloaded_file_path = drive_down(url, directory)
     elif "huggingface.co" in url:
         downloaded_file_path = hf_down(url, directory, hf_token, romanize)
-    elif "civitai.com" in url:
+    elif "civitai." in url:
+        url = url.replace("civitai.red", "civitai.com")
         if not civitai_api_key:
             msg = "You need an API key to download Civitai models."
             print(f"\033[91m{msg}\033[0m")
@@ -393,7 +405,10 @@ def download_things(directory, url, hf_token="", civitai_api_key="", romanize=Fa
             gr.Warning(msg)
             downloaded_file_path = civ_api_down(url, directory, civitai_api_key, civ_filename)
     else:
-        os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
+        os.system(
+            f'wget -c -nv '
+            f'-P "{directory}" "{url}"'
+        )
 
     return downloaded_file_path