Skip to content

C
coderai
Commit f606b9a7 authored by Your Name's avatar Your Name
Browse files
Options

Add multi-cache support for cached model commands 

- Add get_all_cache_dirs() to find GGUF, HuggingFace, and Diffusers caches
- Update --list-cached-models to show all cache locations
- Update --remove-all-models to clean all cache directories
- Update --remove-model to search across all caches
- Add better error handling for diffusers image extraction
parent ab253a98
Hide whitespace changes
Inline Side-by-side
Showing with 148 additions and 54 deletions
coderai
View file @ f606b9a7
......@@ -116,6 +116,28 @@ def get_model_cache_dir() -> str:
pathlib.Path(cache_dir).mkdir(parents=True, exist_ok=True)
return cache_dir
def get_all_cache_dirs() -> dict:
"""Get all model cache directories."""
caches = {}
cache_home = os.environ.get('XDG_CACHE_HOME', os.path.expanduser('~/.cache'))
# Coderai GGUF cache
coderai_cache = os.path.join(cache_home, 'coderai', 'models')
if os.path.exists(coderai_cache):
caches['coderai'] = coderai_cache
# HuggingFace cache (for .safetensors, PyTorch models, etc.)
hf_cache = os.path.join(cache_home, 'huggingface')
if os.path.exists(hf_cache):
caches['huggingface'] = hf_cache
# Local diffusers cache (often stored locally by apps)
local_diffusers = os.path.expanduser('~/.cache/diffusers')
if os.path.exists(local_diffusers):
caches['diffusers'] = local_diffusers
return caches
def get_cached_model_path(url: str) -> Optional[str]:
"""Check if a model URL is already cached. Returns path if cached, None otherwise."""
cache_dir = get_model_cache_dir()
......@@ -3767,7 +3789,16 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
# Extract images
images = []
for img in result.images:
try:
result_images = result.images
except Exception as img_err:
print(f"Warning: Could not access result.images: {img_err}")
# Try alternative: result might have 'image' or 'output'
result_images = getattr(result, 'image', None) or getattr(result, 'output', None)
if result_images is None:
raise Exception(f"Could not extract images from diffusers result: {img_err}")
for img in result_images:
# Convert to base64
import base64
import io
......@@ -5212,88 +5243,151 @@ def main():
# Handle --list-cached-models
if args.list_cached_models:
print("\nListing cached models...")
cache_dir = get_model_cache_dir()
print(f"Cache directory: {cache_dir}")
print("")
print("\n=== Listing Cached Models ===")
if not os.path.exists(cache_dir):
print("Cache directory does not exist.")
caches = get_all_cache_dirs()
if not caches:
print("No model cache directories found.")
sys.exit(0)
files = os.listdir(cache_dir)
if not files:
print("No cached models found.")
all_files = []
for cache_name, cache_dir in caches.items():
print(f"\n--- {cache_name.upper()} Cache ({cache_dir}) ---")
if not os.path.exists(cache_dir):
print(f" (directory does not exist)")
continue
files = os.listdir(cache_dir)
if not files:
print(f" No cached files.")
continue
# For diffusers, show directory structure
if cache_name == 'diffusers':
for root, dirs, files in os.walk(cache_dir):
for f in files:
filepath = os.path.join(root, f)
rel_path = os.path.relpath(filepath, cache_dir)
size = os.path.getsize(filepath)
all_files.append((cache_name, rel_path, size))
else:
for f in files:
filepath = os.path.join(cache_dir, f)
if os.path.isfile(filepath):
size = os.path.getsize(filepath)
all_files.append((cache_name, f, size))
if not all_files:
print("\nNo cached models found.")
sys.exit(0)
print(f"Found {len(files)} cached files:")
total_size = 0
for f in sorted(files):
filepath = os.path.join(cache_dir, f)
size = os.path.getsize(filepath)
total_size += size
size_mb = size / (1024 * 1024)
print(f" {f} ({size_mb:.1f} MB)")
# Calculate totals
total_size = sum(size for _, _, size in all_files)
print(f"\n=== Summary ===")
print(f"Total: {len(all_files)} files, {total_size / (1024*1024*1024):.2f} GB")
print("\nCache locations:")
for cache_name, cache_dir in caches.items():
print(f" {cache_name}: {cache_dir}")
print(f"\nTotal: {len(files)} files, {total_size / (1024*1024*1024):.2f} GB")
sys.exit(0)
# Handle --remove-all-models
if args.remove_all_models:
print("\nRemoving all cached models...")
cache_dir = get_model_cache_dir()
print("\n=== Removing All Cached Models ===")
if not os.path.exists(cache_dir):
print("Cache directory does not exist.")
sys.exit(0)
import shutil
caches = get_all_cache_dirs()
files = os.listdir(cache_dir)
if not files:
print("No cached models to remove.")
if not caches:
print("No cache directories found.")
sys.exit(0)
print(f"Found {len(files)} cached files. Deleting...")
import shutil
for f in files:
filepath = os.path.join(cache_dir, f)
os.remove(filepath)
print(f" Deleted: {f}")
total_removed = 0
for cache_name, cache_dir in caches.items():
if not os.path.exists(cache_dir):
continue
files = os.listdir(cache_dir)
if not files:
continue
print(f"\nRemoving from {cache_name} cache ({cache_dir})...")
print(f" Found {len(files)} file(s). Deleting...")
# For diffusers, remove entire directory tree
if cache_name == 'diffusers':
for item in os.listdir(cache_dir):
item_path = os.path.join(cache_dir, item)
if os.path.isdir(item_path):
shutil.rmtree(item_path)
else:
os.remove(item_path)
print(f" Deleted: {item}")
total_removed += 1
else:
for f in files:
filepath = os.path.join(cache_dir, f)
os.remove(filepath)
print(f" Deleted: {f}")
total_removed += 1
print(f"\nAll {len(files)} cached models removed.")
print(f"\n=== Removed {total_removed} item(s) from all caches ===")
sys.exit(0)
# Handle --remove-model
if args.remove_model:
print(f"\nRemoving cached model matching: {args.remove_model}")
cache_dir = get_model_cache_dir()
print(f"\n=== Removing Cached Model Matching: {args.remove_model} ===")
if not os.path.exists(cache_dir):
print("Cache directory does not exist.")
sys.exit(0)
import shutil
caches = get_all_cache_dirs()
files = os.listdir(cache_dir)
# Find files that contain the search term
matching = [f for f in files if args.remove_model.lower() in f.lower()]
if not caches:
print("No cache directories found.")
sys.exit(0)
if not matching:
all_matching = []
for cache_name, cache_dir in caches.items():
if not os.path.exists(cache_dir):
continue
# For diffusers, search in subdirectories
if cache_name == 'diffusers':
for root, dirs, files in os.walk(cache_dir):
for f in files:
if args.remove_model.lower() in f.lower():
filepath = os.path.join(root, f)
rel_path = os.path.relpath(filepath, cache_dir)
size = os.path.getsize(filepath)
all_matching.append((cache_name, rel_path, filepath, size))
else:
files = os.listdir(cache_dir)
for f in files:
if args.remove_model.lower() in f.lower():
filepath = os.path.join(cache_dir, f)
if os.path.isfile(filepath):
size = os.path.getsize(filepath)
all_matching.append((cache_name, f, filepath, size))
if not all_matching:
print(f"No cached models found matching: {args.remove_model}")
print(f"\nUse --list-cached-models to see available models.")
sys.exit(0)
print(f"Found {len(matching)} matching file(s):")
for f in matching:
filepath = os.path.join(cache_dir, f)
size = os.path.getsize(filepath)
print(f" {f} ({size / (1024*1024):.1f} MB)")
print(f"\nFound {len(all_matching)} matching file(s):")
for cache_name, filename, filepath, size in all_matching:
print(f" [{cache_name}] {filename} ({size / (1024*1024):.1f} MB)")
# Confirm before deleting
print(f"\nDeleting {len(matching)} file(s)...")
for f in matching:
filepath = os.path.join(cache_dir, f)
os.remove(filepath)
print(f" Deleted: {f}")
print(f"\nDeleting {len(all_matching)} file(s)...")
for cache_name, filename, filepath, size in all_matching:
try:
os.remove(filepath)
print(f" Deleted: [{cache_name}] {filename}")
except Exception as e:
print(f" Failed to delete {filename}: {e}")
print(f"\nRemoved {len(matching)} cached model(s).")
print(f"\nRemoved {len(all_matching)} cached model file(s).")
sys.exit(0)
# Get model names from args - support multiple models
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment