Upload modutils.py

modutils.py

@@ -31,6 +31,101 @@ from env import (HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,
     HF_MODEL_USER_EX, HF_MODEL_USER_LIKES, DIFFUSERS_FORMAT_LORAS,
     DIRECTORY_LORAS, HF_READ_TOKEN, HF_TOKEN, CIVITAI_API_KEY)
 
+OUTPUT_CACHE_DIR = Path(os.getenv("OUTPUT_CACHE_DIR", "outputs"))
+OUTPUT_CACHE_MAX_FILES = max(16, int(os.getenv("OUTPUT_CACHE_MAX_FILES", "256")))
+OUTPUT_CACHE_MAX_BYTES = max(512 * 1024**2, int(float(os.getenv("OUTPUT_CACHE_MAX_GB", "4")) * 1024**3))
+CIVITAI_LORA_CACHE_MAX_FILES = max(16, int(os.getenv("CIVITAI_LORA_CACHE_MAX_FILES", "128")))
+CIVITAI_LORA_CACHE_MAX_BYTES = max(4 * 1024**3, int(float(os.getenv("CIVITAI_LORA_CACHE_MAX_GB", "32")) * 1024**3))
+CIVITAI_LORA_CACHE_INDEX = Path(DIRECTORY_LORAS) / ".civitai_lora_lru.json"
+CIVITAI_ALLOWED_LORA_BASE_MODELS = ['Flux.1 D', 'Flux.1 S', 'Flux.1 Kontext', 'Flux.1 Krea', 'Flux.2 D', 'Flux.2 Klein 4B-base', 'Flux.2 Klein 9B', 'Flux.2 Klein 9B-base', 'SD 1.5', 'SD 1.5 Hyper', 'SD 1.5 LCM', 'SDXL 0.9', 'SDXL 1.0', 'SDXL Hyper', 'SDXL Lightning', 'SDXL 1.0 LCM', 'Pony', 'Illustrious', 'NoobAI']
+
+
+def _prune_paths_lru(paths, max_files: int, max_bytes: int, protect=None):
+    protect = {str(Path(p).resolve()) for p in (protect or []) if p}
+    entries = []
+    for raw in paths:
+        try:
+            p = Path(raw)
+            if not p.is_file():
+                continue
+            st = p.stat()
+            entries.append((p, int(st.st_size), float(st.st_mtime)))
+        except Exception:
+            continue
+    total = sum(size for _, size, _ in entries)
+    entries.sort(key=lambda item: item[2])
+    while entries and (len(entries) > max_files or total > max_bytes):
+        victim, size, _ = entries.pop(0)
+        if str(victim.resolve()) in protect:
+            entries.append((victim, size, float("inf")))
+            entries.sort(key=lambda item: item[2])
+            if all(str(p.resolve()) in protect for p, _, _ in entries):
+                break
+            continue
+        try:
+            victim.unlink()
+            total -= size
+            print(f"[cache] pruned {victim}")
+        except Exception as e:
+            print(f"[cache] prune failed {victim} {type(e).__name__}: {e}")
+    return total
+
+
+def _prune_generated_outputs(protect=None):
+    OUTPUT_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    _prune_paths_lru(OUTPUT_CACHE_DIR.glob("*.png"), OUTPUT_CACHE_MAX_FILES, OUTPUT_CACHE_MAX_BYTES, protect=protect)
+    try:
+        return {str(path.resolve()) for path in OUTPUT_CACHE_DIR.glob("*.png") if path.is_file()}
+    except Exception:
+        return set()
+
+
+def _load_civitai_lora_index():
+    try:
+        data = json.loads(CIVITAI_LORA_CACHE_INDEX.read_text(encoding="utf-8"))
+        return data if isinstance(data, dict) else {}
+    except Exception:
+        return {}
+
+
+def _save_civitai_lora_index(data):
+    try:
+        CIVITAI_LORA_CACHE_INDEX.parent.mkdir(parents=True, exist_ok=True)
+        temp = CIVITAI_LORA_CACHE_INDEX.with_suffix(".tmp")
+        temp.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
+        os.replace(temp, CIVITAI_LORA_CACHE_INDEX)
+    except Exception as e:
+        print(f"[civitai] lora cache index save failed: {type(e).__name__}: {e}")
+
+
+def _record_civitai_lora_cache(path: str):
+    try:
+        current = Path(path).resolve()
+        lora_root = Path(DIRECTORY_LORAS).resolve()
+        if not current.is_file() or current.parent != lora_root:
+            return
+        data = _load_civitai_lora_index()
+        normalized = {}
+        for raw, touched in data.items():
+            try:
+                p = Path(raw).resolve()
+                if p.is_file() and p.parent == lora_root:
+                    normalized[str(p)] = float(touched)
+            except Exception:
+                pass
+        normalized[str(current)] = time.time()
+        paths = list(normalized.keys())
+        _prune_paths_lru(paths, CIVITAI_LORA_CACHE_MAX_FILES, CIVITAI_LORA_CACHE_MAX_BYTES, protect=[str(current)])
+        normalized = {raw: ts for raw, ts in normalized.items() if Path(raw).is_file()}
+        _save_civitai_lora_index(normalized)
+    except Exception as e:
+        print(f"[civitai] lora cache accounting failed: {type(e).__name__}: {e}")
+
+
+def is_allowed_civitai_lora_base_model(base_model: str) -> bool:
+    value = str(base_model or "").strip()
+    return value in CIVITAI_ALLOWED_LORA_BASE_MODELS
+
 MODEL_TYPE_DICT = {
     "diffusers:StableDiffusionPipeline": "SD 1.5",
     "diffusers:StableDiffusionXLPipeline": "SDXL",
@@ -339,14 +434,20 @@ def get_civitai_active_api_origin(force_refresh: bool = False, api_key: str = ""
     if CIVITAI_ACTIVE_API_ORIGIN and not force_refresh:
         return CIVITAI_ACTIVE_API_ORIGIN
     session = create_retry_session(total=2, backoff_factor=0.5)
-    for origin in CIVITAI_API_ORIGIN_CANDIDATES:
-        if probe_civitai_api_origin(session, origin, api_key=api_key):
-            set_civitai_active_api_origin(origin)
-            print(f"[civitai] selected api origin: {CIVITAI_ACTIVE_API_ORIGIN}")
-            return CIVITAI_ACTIVE_API_ORIGIN
-    set_civitai_active_api_origin(CIVITAI_DEFAULT_ORIGIN)
-    print(f"[civitai] api probe fallback origin: {CIVITAI_ACTIVE_API_ORIGIN}")
-    return CIVITAI_ACTIVE_API_ORIGIN
+    try:
+        for origin in CIVITAI_API_ORIGIN_CANDIDATES:
+            if probe_civitai_api_origin(session, origin, api_key=api_key):
+                set_civitai_active_api_origin(origin)
+                print(f"[civitai] selected api origin: {CIVITAI_ACTIVE_API_ORIGIN}")
+                return CIVITAI_ACTIVE_API_ORIGIN
+        set_civitai_active_api_origin(CIVITAI_DEFAULT_ORIGIN)
+        print(f"[civitai] api probe fallback origin: {CIVITAI_ACTIVE_API_ORIGIN}")
+        return CIVITAI_ACTIVE_API_ORIGIN
+    finally:
+        try:
+            session.close()
+        except Exception:
+            pass
 
 def get_civitai_active_api_base(force_refresh: bool = False, api_key: str = ""):
     if CIVITAI_ACTIVE_API_BASE and not force_refresh:
@@ -622,6 +723,7 @@ def resolve_civitai_model_page_to_download_url(url: str, api_key: str = ""):
     headers = get_civitai_headers(api_key if parts.netloc.lower().endswith("civitai.com") else "")
     headers['Referer'] = f"{parts.scheme or 'https'}://{parts.netloc}/"
     session = create_retry_session(total=CIVITAI_RESOLVE_RETRY_TOTAL, backoff_factor=CIVITAI_RESOLVE_RETRY_BACKOFF)
+    r = None
     try:
         r = session.get(raw, headers=headers, timeout=CIVITAI_RESOLVE_TIMEOUT)
         if not r.ok:
@@ -638,6 +740,16 @@ def resolve_civitai_model_page_to_download_url(url: str, api_key: str = ""):
     except Exception as e:
         print(f"Failed to resolve Civitai model page URL: {sanitize_url_for_log(raw)} {type(e).__name__}: {sanitize_sensitive_log_text(e)}")
         return raw
+    finally:
+        try:
+            if r is not None:
+                r.close()
+        except Exception:
+            pass
+        try:
+            session.close()
+        except Exception:
+            pass
 
 def normalize_civitai_input_url(url: str, api_key: str = ""):
     raw = str(url or "").strip()
@@ -679,8 +791,9 @@ def resolve_civitai_download_url(url: str, civitai_api_key: str = "", max_tries:
     last_error = None
     for attempt in range(1, max_tries + 1):
         response = None
+        session = create_retry_session(total=3, backoff_factor=1.0)
         try:
-            response = create_retry_session(total=3, backoff_factor=1.0).get(
+            response = session.get(
                 dl_url,
                 headers=headers,
                 allow_redirects=False,
@@ -712,6 +825,10 @@ def resolve_civitai_download_url(url: str, civitai_api_key: str = "", max_tries:
                     response.close()
             except Exception:
                 pass
+            try:
+                session.close()
+            except Exception:
+                pass
         if attempt < max_tries:
             time.sleep(min(3.0, 0.8 * attempt))
     if last_error is not None:
@@ -819,6 +936,7 @@ def request_json_data(url, api_key: str = ""):
         if attempt > 1:
             headers["Connection"] = "close"
         endpoint_url = ""
+        result = None
         try:
             json_data, endpoint_url, result = request_civitai_api_json(
                 endpoint_path,
@@ -849,6 +967,11 @@ def request_json_data(url, api_key: str = ""):
                 f"error={type(e).__name__}: {sanitize_sensitive_log_text(e)}"
             )
         finally:
+            try:
+                if result is not None:
+                    result.close()
+            except Exception:
+                pass
             try:
                 session.close()
             except Exception:
@@ -874,6 +997,7 @@ class ModelInformation:
             self.description = ""
         self.model_name = json_data.get("model", {}).get("name", "")
         self.model_type = json_data.get("model", {}).get("type", "")
+        self.base_model = json_data.get("baseModel", "")
         self.nsfw = json_data.get("model", {}).get("nsfw", False)
         self.poi = json_data.get("model", {}).get("poi", False)
         self.images = [img.get("url", "") for img in json_data.get("images", [])]
@@ -1145,6 +1269,13 @@ def download_things(directory, url, hf_token="", civitai_api_key="", romanize=Fa
             print(f"Civitai download URL normalized: {sanitize_url_for_log(url)} -> {sanitize_url_for_log(normalized_url)}")
         model_profile = retrieve_model_info(normalized_url, api_key=civitai_api_key)
         selected_file = model_profile.selected_file if model_profile else {}
+        if Path(directory).name == Path(DIRECTORY_LORAS).name and model_profile:
+            if str(model_profile.model_type or "").upper() not in {"LORA", "LOCON", "DORA"}:
+                print(f"[civitai] rejected non-LoRA model type={model_profile.model_type!r}")
+                return None
+            if not is_allowed_civitai_lora_base_model(model_profile.base_model):
+                print(f"[civitai] rejected LoRA base model={model_profile.base_model!r}")
+                return None
         if model_profile and model_profile.download_url:
             url = model_profile.download_url
             filename = model_profile.filename_url or ""
@@ -1333,15 +1464,20 @@ def save_gallery_images(images, model_name="", progress=gr.Progress(track_tqdm=T
     output_images = []
     output_paths = []
     for i, image in enumerate(images):
+        OUTPUT_CACHE_DIR.mkdir(parents=True, exist_ok=True)
         filename = f"{basename}{str(i + 1)}.png"
+        target_path = (OUTPUT_CACHE_DIR / filename).resolve()
         oldpath = Path(image[0])
         newpath = oldpath.resolve() if oldpath.exists() else oldpath
         try:
             if oldpath.exists():
                 source_path = oldpath.resolve()
-                target_path = Path(filename).resolve()
                 if source_path != target_path:
-                    shutil.copy2(str(source_path), str(target_path))
+                    owned_temp = source_path.parent == Path(tempfile.gettempdir()).resolve() and source_path.name.startswith("modimg_")
+                    if owned_temp:
+                        shutil.move(str(source_path), str(target_path))
+                    else:
+                        shutil.copy2(str(source_path), str(target_path))
                     newpath = target_path
                 else:
                     newpath = source_path
@@ -1351,6 +1487,7 @@ def save_gallery_images(images, model_name="", progress=gr.Progress(track_tqdm=T
         finally:
             output_paths.append(str(newpath))
             output_images.append((str(newpath), str(filename)))
+    _prune_generated_outputs(protect=output_paths)
     progress(1, desc="Gallery updated.")
     return gr.update(value=output_images), gr.update(value=output_paths, visible=True)
 
@@ -1360,6 +1497,10 @@ def save_gallery_history(images, files, history_gallery, history_files, progress
     if not history_files: history_files = []
     output_gallery = images + history_gallery
     output_files = files + history_files
+    survivors = _prune_generated_outputs(protect=files)
+    if survivors:
+        output_files = [path for path in output_files if str(Path(path).resolve()) in survivors]
+        output_gallery = [item for item in output_gallery if item and str(Path(item[0]).resolve()) in survivors]
     return gr.update(value=output_gallery), gr.update(value=output_files, visible=True)
 
 def save_image_history(image, gallery, files, model_name: str, progress=gr.Progress(track_tqdm=True)):
@@ -1369,7 +1510,9 @@ def save_image_history(image, gallery, files, model_name: str, progress=gr.Progr
     try:
         basename = f"{model_name.split('/')[-1]}_{datetime.now(FILENAME_TIMEZONE).strftime('%Y%m%d_%H%M%S')}"
         if image is None or not isinstance(image, (str, Image.Image, np.ndarray, tuple)): return gr.update(), gr.update()
+        OUTPUT_CACHE_DIR.mkdir(parents=True, exist_ok=True)
         filename = f"{basename}.png"
+        target_path = (OUTPUT_CACHE_DIR / filename).resolve()
         if isinstance(image, tuple): image = image[0]
         if isinstance(image, str):
             oldpath = image
@@ -1384,10 +1527,21 @@ def save_image_history(image, gallery, files, model_name: str, progress=gr.Progr
         oldpath = Path(oldpath)
         newpath = oldpath
         if oldpath.exists():
-            shutil.copy2(str(oldpath.resolve()), str(Path(filename).resolve()))
-            newpath = Path(filename).resolve()
+            source_path = oldpath.resolve()
+            owned_temp = source_path.parent == Path(tempfile.gettempdir()).resolve() and source_path.name.startswith(("modimg_", "history_"))
+            if source_path != target_path:
+                if owned_temp:
+                    shutil.move(str(source_path), str(target_path))
+                else:
+                    shutil.copy2(str(source_path), str(target_path))
+            newpath = target_path
         files.insert(0, str(newpath))
         gallery.insert(0, (str(newpath), str(filename)))
+        survivors = _prune_generated_outputs(protect=[str(newpath)])
+        if survivors:
+            kept = [(g, f) for g, f in zip(gallery, files) if str(Path(f).resolve()) in survivors]
+            gallery = [g for g, _ in kept]
+            files = [f for _, f in kept]
     except Exception as e:
         log_error(e)
     finally:
@@ -1670,6 +1824,8 @@ def finalize_downloaded_lora_path(file_path: str, source_url: str = ""):
             if normalized_url:
                 loras_url_to_path_dict[normalized_url] = final_path
     update_lora_dict(final_path)
+    if source_url and is_civitai_url(source_url):
+        _record_civitai_lora_cache(final_path)
     return final_path
 
 def download_lora(dl_urls: str):
@@ -2027,6 +2183,7 @@ def get_civitai_info(path):
     headers = get_civitai_headers(CIVITAI_API_KEY)
     endpoint_path = '/model-versions/by-hash/'
     session = create_retry_session()
+    r = None
 
     import hashlib
     sha256_hash = hashlib.sha256()
@@ -2047,6 +2204,16 @@ def get_civitai_info(path):
     except Exception as e:
         print(f"Civitai by-hash lookup failed: {path} {type(e).__name__}: {e}")
         return default
+    finally:
+        try:
+            if r is not None:
+                r.close()
+        except Exception:
+            pass
+        try:
+            session.close()
+        except Exception:
+            pass
     if not r.ok:
         print(f"Civitai by-hash lookup status={r.status_code}: {path}")
         if r.status_code == 404:
@@ -2115,7 +2282,7 @@ def build_civitai_choice_name(item: dict) -> str:
     base_model_name = "Pony🐴" if item.get('base_model') == "Pony" else item.get('base_model', '')
     return f"{item.get('name', '')} (for {base_model_name} / By: {item.get('creator', '')} / Tags: {', '.join(item.get('tags', []))})"
 
-def search_lora_on_civitai(query: str, allow_model: list[str] = ["Pony", "SDXL 1.0"], limit: int = 100,
+def search_lora_on_civitai(query: str, allow_model: list[str] | None = None, limit: int = 100,
                            sort: str = "Highest Rated", period: str = "AllTime", tag: str = "", user: str = "", page: int = 1):
     headers = get_civitai_headers(CIVITAI_API_KEY)
     endpoint_path = '/models'
@@ -2127,6 +2294,7 @@ def search_lora_on_civitai(query: str, allow_model: list[str] = ["Pony", "SDXL 1
     if user:
         params["username"] = user
     session = create_retry_session()
+    r = None
     try:
         json, _, r = request_civitai_api_json(
             endpoint_path,
@@ -2140,6 +2308,16 @@ def search_lora_on_civitai(query: str, allow_model: list[str] = ["Pony", "SDXL 1
     except Exception as e:
         print(f"Civitai search failed: query={query!r} page={page} {type(e).__name__}: {e}")
         return None
+    finally:
+        try:
+            if r is not None:
+                r.close()
+        except Exception:
+            pass
+        try:
+            session.close()
+        except Exception:
+            pass
     if not r.ok or not json:
         print(f"Civitai search status={r.status_code}: query={query!r} page={page}")
         return None
@@ -2147,7 +2325,7 @@ def search_lora_on_civitai(query: str, allow_model: list[str] = ["Pony", "SDXL 1
         print(f"Civitai search returned no items key: query={query!r} page={page}")
         return None
     items = []
-    allowed_models = set(allow_model)
+    allowed_models = set(allow_model or CIVITAI_ALLOWED_LORA_BASE_MODELS)
     for j in json['items']:
         model_versions = j.get('modelVersions') if isinstance(j, dict) else []
         if not isinstance(model_versions, list):
@@ -2163,13 +2341,7 @@ def search_lora_on_civitai(query: str, allow_model: list[str] = ["Pony", "SDXL 1
 
 CIVITAI_SORT = ["Highest Rated", "Most Downloaded", "Most Liked", "Most Discussed", "Most Collected", "Most Buzz", "Newest"]
 CIVITAI_PERIOD = ["AllTime", "Year", "Month", "Week", "Day"]
-CIVITAI_BASEMODEL_DEFAULT = ["Chroma", "Flux.1 D", "Flux.1 S", "Flux.1 Kontext", "HiDream", "Hunyuan Video",
-                             "Illustrious", "NoobAI", "Other", "Pony", "SD 1.4", "SD 1.5", "SD 1.5 Hyper",
-                             "SD 1.5 LCM", "SD 2.0", "SD 2.1", "SD 2.1 768", "SDXL 0.9", "SDXL 1.0", "SDXL Hyper",
-                             "SDXL Lightning", "Wan Video", "Anima", "Flux.1 Krea", "Flux.2 D", "Flux.2 Klein 4B-base",
-                             "Flux.2 Klein 9B", "Flux.2 Klein 9B-base", "Grok", "LTXV 2.3", "LTXV2", "Qwen", "SDXL 1.0 LCM",
-                             "Wan Video 1.3B t2v", "Wan Video 14B i2v 480p", "Wan Video 14B i2v 720p", "Wan Video 14B t2v",
-                             "Wan Video 2.2 I2V-A14B", "Wan Video 2.2 T2V-A14B", "Wan Video 2.2 TI2V-5B", "ZImageBase", "ZImageTurbo"]
+CIVITAI_BASEMODEL_DEFAULT = ['Flux.1 D', 'Flux.1 S', 'Flux.1 Kontext', 'Flux.1 Krea', 'Flux.2 D', 'Flux.2 Klein 4B-base', 'Flux.2 Klein 9B', 'Flux.2 Klein 9B-base', 'SD 1.5', 'SD 1.5 Hyper', 'SD 1.5 LCM', 'SDXL 0.9', 'SDXL 1.0', 'SDXL Hyper', 'SDXL Lightning', 'SDXL 1.0 LCM', 'Pony', 'Illustrious', 'NoobAI']
 CIVITAI_BASEMODEL = CIVITAI_BASEMODEL_DEFAULT.copy()
 
 
@@ -2247,6 +2419,7 @@ def get_civitai_tag():
     headers = {'User-Agent': user_agent, 'content-type': 'application/json'}
     params = {'limit': 200}
     session = create_retry_session()
+    r = None
     try:
         json_data, _, r = request_civitai_api_json(
             '/tags',
@@ -2271,6 +2444,16 @@ def get_civitai_tag():
     except Exception as e:
         log_warning(e)
         return default
+    finally:
+        try:
+            if r is not None:
+                r.close()
+        except Exception:
+            pass
+        try:
+            session.close()
+        except Exception:
+            pass
 
 LORA_BASE_MODEL_DICT = {
     "diffusers:StableDiffusionPipeline": ["SD 1.5"],
@@ -2706,9 +2889,9 @@ def read_safetensors_key(path: str):
     except Exception as e:
         log_error(e)
     finally:
+        gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
-        gc.collect()
     return keys
 
 def get_model_type_from_key(path: str):