Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import gc
+import re
 import gradio as gr
 import numpy as np
 import spaces
@@ -8,7 +9,7 @@ import random
 from PIL import Image
 from typing import Iterable, Optional
 
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, HfApi
 from safetensors.torch import load_file as safetensors_load_file
 
 from gradio.themes import Soft
@@ -119,17 +120,119 @@ from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
 dtype = torch.bfloat16
 
+# ------------------------------------------------------------
+# AIO versioning
+# ------------------------------------------------------------
+AIO_REPO_ID = "Pr0f3ssi0n4ln00b/Phr00t-Qwen-Rapid-AIO"
+DEFAULT_AIO_VERSION = "v19"
+_VERSION_RE = re.compile(r"^v\d+$")
+
+def discover_aio_versions(repo_id: str) -> list[str]:
+    """
+    Discovers versions that follow vXX/transformer/ in the HF repo.
+    Returns sorted list like: ['v19', 'v21', ...]
+    """
+    api = HfApi()
+    try:
+        files = api.list_repo_files(repo_id=repo_id, repo_type="model")
+    except Exception as e:
+        print(f"⚠️ Could not list repo files for {repo_id}: {e}")
+        return [DEFAULT_AIO_VERSION]
+
+    versions = set()
+    for p in files:
+        if "/transformer/" not in p:
+            continue
+        head = p.split("/transformer/", 1)[0]  # "v19"
+        if _VERSION_RE.fullmatch(head):
+            versions.add(head)
+
+    if not versions:
+        versions = {DEFAULT_AIO_VERSION}
+
+    return sorted(versions, key=lambda x: int(x[1:]))
+
+AVAILABLE_AIO_VERSIONS = discover_aio_versions(AIO_REPO_ID)
+
+# Track currently loaded transformer version
+CURRENT_AIO_VERSION: Optional[str] = None
+
+def _free_cuda():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+@spaces.GPU
+def switch_aio_version(version: str):
+    """
+    Loads transformer weights from {version}/transformer/ into the already-created pipeline.
+    """
+    global CURRENT_AIO_VERSION, pipe
+
+    if version is None or str(version).strip() == "":
+        version = DEFAULT_AIO_VERSION
+
+    if CURRENT_AIO_VERSION == version:
+        return gr.update(value=f"✅ Already using {version}")
+
+    _free_cuda()
+
+    subfolder = f"{version}/transformer"
+    print(f"🔁 Switching AIO transformer to: {AIO_REPO_ID} / {subfolder}")
+
+    old_transformer = getattr(pipe, "transformer", None)
+
+    new_transformer = QwenImageTransformer2DModel.from_pretrained(
+        AIO_REPO_ID,
+        subfolder=subfolder,
+        torch_dtype=dtype,
+        device_map="cuda",
+    )
+    pipe.transformer = new_transformer
+
+    # Re-apply FA3 Optimization
+    try:
+        pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+        print("Flash Attention 3 Processor set successfully.")
+    except Exception as e:
+        print(f"Warning: Could not set FA3 processor: {e}")
+
+    # Best-effort free old transformer reference
+    try:
+        del old_transformer
+    except Exception:
+        pass
+
+    _free_cuda()
+
+    CURRENT_AIO_VERSION = version
+    return gr.update(value=f"✅ Loaded {version} ({subfolder}/)")
+
+def refresh_aio_versions():
+    global AVAILABLE_AIO_VERSIONS
+    AVAILABLE_AIO_VERSIONS = discover_aio_versions(AIO_REPO_ID)
+    new_default = DEFAULT_AIO_VERSION if DEFAULT_AIO_VERSION in AVAILABLE_AIO_VERSIONS else AVAILABLE_AIO_VERSIONS[0]
+    return (
+        gr.update(choices=AVAILABLE_AIO_VERSIONS, value=new_default),
+        gr.update(value=f"🔄 Found: {', '.join(AVAILABLE_AIO_VERSIONS)}")
+    )
+
+# ------------------------------------------------------------
+# Create pipeline (loads DEFAULT_AIO_VERSION only)
+# ------------------------------------------------------------
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "Qwen/Qwen-Image-Edit-2511",
     transformer=QwenImageTransformer2DModel.from_pretrained(
-        "Pr0f3ssi0n4ln00b/Phr00t-Qwen-Rapid-AIO",
-        subfolder="transformer",
+        AIO_REPO_ID,
+        subfolder=f"{DEFAULT_AIO_VERSION}/transformer",
         torch_dtype=dtype,
         device_map="cuda",
     ),
     torch_dtype=dtype,
 ).to(device)
 
+CURRENT_AIO_VERSION = DEFAULT_AIO_VERSION
+
 # Apply FA3 Optimization
 try:
     pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
@@ -285,11 +388,9 @@ LOADED_ADAPTERS = set()
 # Helpers: resolution
 # ============================================================
 
-
 def _round8(x: int) -> int:
     return max(8, (int(x) // 8) * 8)
 
-
 def compute_dimensions(image: Image.Image, long_edge: int) -> tuple[int, int]:
     w, h = image.size
     if w >= h:
@@ -300,25 +401,20 @@ def compute_dimensions(image: Image.Image, long_edge: int) -> tuple[int, int]:
         new_w = int(round(long_edge * (w / h)))
     return _round8(new_w), _round8(new_h)
 
-
 def get_target_long_edge_for_lora(lora_adapter: str) -> int:
     spec = ADAPTER_SPECS.get(lora_adapter, {})
     return int(spec.get("target_long_edge", 1024))
 
-
 # ============================================================
 # Helpers: multi-input routing + gallery normalization
 # ============================================================
 
-
 def lora_requires_two_images(lora_adapter: str) -> bool:
     return bool(ADAPTER_SPECS.get(lora_adapter, {}).get("requires_two_images", False))
 
-
 def image2_label_for_lora(lora_adapter: str) -> str:
     return str(ADAPTER_SPECS.get(lora_adapter, {}).get("image2_label", "Upload Reference (Image 2)"))
 
-
 def _to_pil_rgb(x) -> Optional[Image.Image]:
     """
     Accepts PIL / numpy / (image, caption) tuples from gr.Gallery and returns PIL RGB.
@@ -345,7 +441,6 @@ def _to_pil_rgb(x) -> Optional[Image.Image]:
     except Exception:
         return None
 
-
 def build_labeled_images(
     img1: Image.Image,
     img2: Optional[Image.Image],
@@ -377,12 +472,10 @@ def build_labeled_images(
 
     return labeled
 
-
 # ============================================================
 # Helpers: BFS alpha key fix
 # ============================================================
 
-
 def _inject_missing_alpha_keys(state_dict: dict) -> dict:
     """
     Diffusers' Qwen LoRA converter expects '<module>.alpha' keys.
@@ -418,7 +511,6 @@ def _inject_missing_alpha_keys(state_dict: dict) -> dict:
 
     return state_dict
 
-
 def _load_lora_weights_with_fallback(repo: str, weight_name: str, adapter_name: str, needs_alpha_fix: bool = False):
     """
     Normal path: pipe.load_lora_weights(repo, weight_name=..., adapter_name=...)
@@ -439,12 +531,10 @@ def _load_lora_weights_with_fallback(repo: str, weight_name: str, adapter_name:
         pipe.load_lora_weights(sd, adapter_name=adapter_name)
         return
 
-
 # ============================================================
 # LoRA loader: single/package + strengths
 # ============================================================
 
-
 def _ensure_loaded_and_get_active_adapters(selected_lora: str):
     spec = ADAPTER_SPECS.get(selected_lora)
     if not spec:
@@ -510,12 +600,10 @@ def _ensure_loaded_and_get_active_adapters(selected_lora: str):
 
     return adapter_names, adapter_weights
 
-
 # ============================================================
 # UI handlers
 # ============================================================
 
-
 def on_lora_change_ui(selected_lora, current_prompt):
     # Preset prompt (fill only if empty)
     if selected_lora != NONE_LORA:
@@ -535,19 +623,18 @@ def on_lora_change_ui(selected_lora, current_prompt):
 
     return prompt_update, img2_update
 
-
 # ============================================================
 # Inference
 # ============================================================
 
-
 @spaces.GPU
 def infer(
     input_image_1,
     input_image_2,
-    input_images_extra,  # NEW: gallery multi-image box
+    input_images_extra,  # gallery multi-image box
     prompt,
     lora_adapter,
+    aio_version,         # NEW: selected AIO version
     seed,
     randomize_seed,
     guidance_scale,
@@ -558,6 +645,10 @@ def infer(
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
+    # Ensure the requested transformer version is loaded
+    if aio_version and aio_version != CURRENT_AIO_VERSION:
+        switch_aio_version(aio_version)
+
     if input_image_1 is None:
         raise gr.Error("Please upload Image 1.")
 
@@ -625,7 +716,6 @@ def infer(
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-
 @spaces.GPU
 def infer_example(input_image, prompt, lora_adapter):
     if input_image is None:
@@ -634,10 +724,20 @@ def infer_example(input_image, prompt, lora_adapter):
     guidance_scale = 1.0
     steps = 4
     # Examples don't supply Image 2 or extra images; and example list doesn't include AnyPose/BFS.
-    result, seed = infer(input_pil, None, None, prompt, lora_adapter, 0, True, guidance_scale, steps)
+    result, seed = infer(
+        input_pil,
+        None,
+        None,
+        prompt,
+        lora_adapter,
+        CURRENT_AIO_VERSION or DEFAULT_AIO_VERSION,  # NEW: keep whatever is loaded
+        0,
+        True,
+        guidance_scale,
+        steps,
+    )
     return result, seed
 
-
 # ============================================================
 # UI
 # ============================================================
@@ -664,7 +764,6 @@ with gr.Blocks() as demo:
                 input_image_1 = gr.Image(label="Upload Image 1 (Base / Target)", type="pil", height=290)
                 input_image_2 = gr.Image(label="Upload Reference (Image 2)", type="pil", height=290, visible=False)
 
-                # NEW: multi-image input box (supports multiple images)
                 input_images_extra = gr.Gallery(
                     label="Upload Additional Images (auto-indexed after Image 1/2)",
                     type="pil",
@@ -685,6 +784,22 @@ with gr.Blocks() as demo:
             with gr.Column():
                 output_image = gr.Image(label="Output Image", interactive=False, format="png", height=353)
 
+                # NEW: AIO version selector + refresh
+                with gr.Row():
+                    aio_version = gr.Dropdown(
+                        label="Phr00t Rapid AIO Version",
+                        choices=AVAILABLE_AIO_VERSIONS,
+                        value=DEFAULT_AIO_VERSION if DEFAULT_AIO_VERSION in AVAILABLE_AIO_VERSIONS else AVAILABLE_AIO_VERSIONS[0],
+                        interactive=True,
+                    )
+                    refresh_versions_btn = gr.Button("↻", scale=0)
+
+                aio_status = gr.Textbox(
+                    label="Model Status",
+                    value=f"Using {CURRENT_AIO_VERSION}",
+                    interactive=False,
+                )
+
                 with gr.Row():
                     lora_choices = [NONE_LORA] + list(ADAPTER_SPECS.keys())
                     lora_adapter = gr.Dropdown(
@@ -706,6 +821,20 @@ with gr.Blocks() as demo:
             outputs=[prompt, input_image_2],
         )
 
+        # On AIO version change: swap transformer
+        aio_version.change(
+            fn=switch_aio_version,
+            inputs=[aio_version],
+            outputs=[aio_status],
+        )
+
+        # Refresh available versions
+        refresh_versions_btn.click(
+            fn=refresh_aio_versions,
+            inputs=[],
+            outputs=[aio_version, aio_status],
+        )
+
         gr.Examples(
             examples=[
                 ["examples/1.jpg", "Transform into anime.", "Photo-to-Anime"],
@@ -746,9 +875,10 @@ with gr.Blocks() as demo:
         inputs=[
             input_image_1,
             input_image_2,
-            input_images_extra,  # NEW
+            input_images_extra,
             prompt,
             lora_adapter,
+            aio_version,      # NEW
             seed,
             randomize_seed,
             guidance_scale,