"""
ColorCraft SD 1.5 - LineArt Specialist Edition
==============================================
SD 1.5 LineArt specialist approach:
- Purpose-built LineArt ControlNet models
- 50% less VRAM usage vs SDXL
- Faster generation times
- Better line art quality (specialized training)
- Storage efficient with smaller models
"""
import os
import gc
import shutil
import gradio as gr
import torch
import numpy as np
from PIL import Image
from huggingface_hub import snapshot_download, hf_hub_download
# Environment setup
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
# Simple compatibility fix at the very beginning
try:
import huggingface_hub
if not hasattr(huggingface_hub, 'cached_download'):
from huggingface_hub import hf_hub_download
huggingface_hub.cached_download = hf_hub_download
print("✅ Applied cached_download compatibility shim")
except Exception as e:
print(f"⚠️ Compatibility shim error: {e}")
# Now import diffusers after the fix
try:
from diffusers import (
StableDiffusionControlNetPipeline, # SD 1.5 pipeline (simpler!)
ControlNetModel,
DDIMScheduler,
EulerAncestralDiscreteScheduler,
StableDiffusionXLControlNetPipeline # SDXL pipeline
)
from diffusers.utils import load_image
print("✅ Diffusers imported successfully")
except ImportError as e:
print(f"❌ Diffusers import error: {e}")
# Try ControlNet preprocessor imports
try:
from controlnet_aux import LineartDetector, LineartAnimeDetector, PidiNetDetector, HEDdetector
print("✅ LineartDetector and soft preprocessors imported successfully")
except ImportError as e:
print(f"❌ Preprocessor import error: {e}")
LineartDetector = None
LineartAnimeDetector = None
PidiNetDetector = None
HEDdetector = None
# Try depth preprocessor (using correct import)
try:
from controlnet_aux import MidasDetector
print("✅ MidasDetector (depth) imported successfully")
DepthEstimator = MidasDetector # Alias for compatibility
except ImportError as e:
print(f"❌ MidasDetector import error: {e}")
DepthEstimator = None
# Fallback depth model
try:
from transformers import DPTImageProcessor, DPTForDepthEstimation
print("✅ DPT depth model imported successfully")
except ImportError as e:
print(f"❌ DPT import error: {e}")
DPTImageProcessor = None
DPTForDepthEstimation = None
import warnings
warnings.filterwarnings("ignore")
# =============================================================================
# STORAGE-EFFICIENT CONFIGURATION
# =============================================================================
SD15_MODEL_ID = "runwayml/stable-diffusion-v1-5"
SDXL_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
# LineArt model options (SD 1.5 ControlNet specialists - using VERIFIED models!)
LINEART_MODELS = {
"standard": {
"type": "controlnet",
"id": "lllyasviel/control_v11p_sd15_lineart", # VERIFIED: Standard LineArt
"preprocessor": "lineart_standard"
},
"anime": {
"type": "controlnet",
"id": "lllyasviel/control_v11p_sd15s2_lineart_anime", # REAL anime ControlNet v1.1!
"preprocessor": "lineart_anime_denoise"
},
"soft_edge": {
"type": "controlnet",
"id": "lllyasviel/control_v11p_sd15_softedge", # Soft edge detection
"preprocessor": "softedge_pidinet"
},
"canny": {
"type": "controlnet",
"id": "lllyasviel/control_v11p_sd15_canny", # VERIFIED: Canny Edge (fallback)
"preprocessor": "canny"
}
}
# SDXL ControlNet options (primary: Union ControlNet)
LINEART_MODELS_SDXL = {
"union": {
"type": "controlnet_sdxl",
"id": "xinsir/controlnet-union-sdxl-1.0",
"preprocessor": "softedge_pidinet"
}
}
DEPTH_MODEL_ID = "lllyasviel/control_v11f1p_sd15_depth"
# Specialized LoRA models for coloring book generation
COLORING_BOOK_LORAS = {
"none": {
"name": "No LoRA",
"path": None,
"trigger": "",
"strength": 0.0
},
"coloring_book_redmond": {
"name": "ColoringBookRedmond V2",
"path": "artificialguybr/ColoringBookRedmond-V2",
"trigger": "ColoringBookAF, coloring book",
"strength": 0.8
},
# Disabled by default (private/unavailable). To enable, set a valid repo and token.
"line_art_simple": {
"name": "Simple LineArt Style (requires auth)",
"path": None,
"trigger": "LineArtAF",
"strength": 0.7
},
"kids_coloring": {
"name": "Kids Coloring Book (requires auth)",
"path": None,
"trigger": "coloring book, simple lines",
"strength": 0.6
}
}
# Cache directories with storage management
CACHE_ROOT = "/data" if os.path.exists("/data") else "."
HF_CACHE_DIR = f"{CACHE_ROOT}/hf_cache"
os.makedirs(HF_CACHE_DIR, exist_ok=True)
# Force all common HF cache env vars to persistent storage
os.environ["HF_HOME"] = HF_CACHE_DIR
os.environ["HUGGINGFACE_HUB_CACHE"] = HF_CACHE_DIR
os.environ["TRANSFORMERS_CACHE"] = HF_CACHE_DIR
os.environ["DIFFUSERS_CACHE"] = HF_CACHE_DIR
HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
# =============================================================================
# STORAGE MANAGEMENT
# =============================================================================
def get_storage_info():
"""Get current storage usage"""
try:
if os.path.exists("/data"):
total, used, free = shutil.disk_usage("/data")
return {
"total_gb": total / (1024**3),
"used_gb": used / (1024**3),
"free_gb": free / (1024**3),
"usage_percent": (used / total) * 100
}
else:
return {
"total_gb": 0,
"used_gb": 0,
"free_gb": 0,
"usage_percent": 0
}
except Exception as e:
print(f"Storage info error: {e}")
return {"total_gb": 0, "used_gb": 0, "free_gb": 0, "usage_percent": 0}
def format_storage_status():
"""Format storage status for display"""
info = get_storage_info()
if info["total_gb"] > 0:
return f"Storage: {info['free_gb']:.1f}GB free of {info['total_gb']:.1f}GB total ({info['usage_percent']:.1f}% used)"
else:
return "Ephemeral storage (no persistent cache)"
def clear_cache_if_needed(min_free_gb=5.0):
"""Clear cache if storage is low - aggressive cleanup for SD 1.5 migration"""
info = get_storage_info()
if info["free_gb"] < min_free_gb:
print(f"⚠️ Low storage: {info['free_gb']:.1f}GB free, clearing cache...")
# Clear HuggingFace cache completely
if os.path.exists(HF_CACHE_DIR):
try:
shutil.rmtree(HF_CACHE_DIR)
os.makedirs(HF_CACHE_DIR, exist_ok=True)
print("✅ HF cache cleared")
except Exception as e:
print(f"Cache clear error: {e}")
# Clear any old SDXL models from custom repo if they exist
custom_cache = "/data"
if os.path.exists(custom_cache):
old_model_patterns = [
"models--stabilityai--stable-diffusion-xl-base-1.0",
"models--diffusers--controlnet-depth-sdxl-1.0",
"models--diffusers--controlnet-canny-sdxl-1.0",
"models--madebyollin--sdxl-vae-fp16-fix",
# SD 1.5 family (remove when migrating to SDXL-only)
"models--runwayml--stable-diffusion-v1-5",
"models--lllyasviel--control_v11p_sd15_lineart",
"models--lllyasviel--control_v11p_sd15s2_lineart_anime",
"models--lllyasviel--control_v11p_sd15_softedge",
"models--lllyasviel--control_v11p_sd15_canny",
"models--lllyasviel--control_v11f1p_sd15_depth",
# Annotators can be large; we keep them elsewhere but clear stale copies under /data
"models--lllyasviel--Annotators",
# Optional depth model if downloaded under /data
"models--Intel--dpt-hybrid-midas"
]
for pattern in old_model_patterns:
for item in os.listdir(custom_cache):
if pattern in item:
try:
path_to_remove = os.path.join(custom_cache, item)
if os.path.isdir(path_to_remove):
shutil.rmtree(path_to_remove)
print(f"🗑️ Cleared old SDXL: {item}")
except Exception as e:
print(f"⚠️ Could not clear {item}: {e}")
# Force garbage collection
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
# Check final storage
final_info = get_storage_info()
print(f"📊 After cleanup: {final_info['free_gb']:.1f}GB free")
return final_info["free_gb"] >= min_free_gb
return True
# =============================================================================
# ADVANCED HF CACHE PRUNING (SDXL-ONLY)
# =============================================================================
def _path_size_gb(path: str) -> float:
try:
total = 0
for root, _, files in os.walk(path):
for f in files:
try:
total += os.path.getsize(os.path.join(root, f))
except Exception:
pass
return total / (1024 ** 3)
except Exception:
return 0.0
def prune_cache_to_sdxl_only(keep_depth: bool = False, keep_lora_repo_ids: list | None = None) -> str:
"""Remove HF cache entries that are not required for SDXL base + Union ControlNet.
Keeps:
- stabilityai/stable-diffusion-xl-base-1.0
- xinsir/controlnet-union-sdxl-1.0
- lllyasviel/Annotators (for PidiNet/LineArt detectors)
- Optional LoRA repos provided in keep_lora_repo_ids
- Optional depth model (Intel/dpt-hybrid-midas) when keep_depth=True
"""
logs: list[str] = []
try:
keep_repos = {
"stabilityai/stable-diffusion-xl-base-1.0",
"xinsir/controlnet-union-sdxl-1.0",
"lllyasviel/Annotators",
}
if keep_lora_repo_ids:
keep_repos.update(keep_lora_repo_ids)
if keep_depth:
keep_repos.add("Intel/dpt-hybrid-midas")
# Prefer official Hugging Face cleanup utilities if available
try:
from huggingface_hub import scan_cache_dir, delete_cache_entries # type: ignore
cache_info = scan_cache_dir(HF_CACHE_DIR)
to_delete = []
for repo in cache_info.repos:
try:
is_model_repo = getattr(repo, "repo_type", "model") == "model"
except Exception:
is_model_repo = True
repo_id = getattr(repo, "repo_id", "")
if is_model_repo and repo_id and repo_id not in keep_repos:
for rev in getattr(repo, "revisions", []):
to_delete.append(rev)
if to_delete:
before_gb = _path_size_gb(HF_CACHE_DIR)
delete_cache_entries(to_delete)
after_gb = _path_size_gb(HF_CACHE_DIR)
logs.append(f"🧹 HF cache vacuumed via API: {before_gb:.2f}GB -> {after_gb:.2f}GB")
else:
logs.append("ℹ️ Nothing to delete via HF cache API")
except Exception as e_api:
logs.append(f"⚠️ HF cache API prune unavailable ({e_api}); using pattern-based cleanup")
# Fallback: delete model directories by pattern under /data and HF cache hub
delete_patterns = [
"models--runwayml--stable-diffusion-v1-5",
"models--lllyasviel--control_v11p_sd15_lineart",
"models--lllyasviel--control_v11p_sd15s2_lineart_anime",
"models--lllyasviel--control_v11p_sd15_softedge",
"models--lllyasviel--control_v11p_sd15_canny",
"models--lllyasviel--control_v11f1p_sd15_depth",
]
if not keep_depth:
delete_patterns.append("models--Intel--dpt-hybrid-midas")
search_roots = [HF_CACHE_DIR, "/data"] if os.path.exists("/data") else [HF_CACHE_DIR]
removed = 0
for root_dir in search_roots:
try:
for item in os.listdir(root_dir):
if any(p in item for p in delete_patterns):
target = os.path.join(root_dir, item)
try:
if os.path.isdir(target):
shutil.rmtree(target)
else:
os.remove(target)
removed += 1
logs.append(f"🗑️ Removed {target}")
except Exception as e_rm:
logs.append(f"⚠️ Could not remove {target}: {e_rm}")
except Exception:
pass
logs.append(f"✅ Pattern-based cleanup complete. Removed {removed} entries")
# Final GC + CUDA cache release
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
info = get_storage_info()
logs.append(f"💾 Free space: {info['free_gb']:.2f}GB of {info['total_gb']:.2f}GB")
except Exception as e:
logs.append(f"❌ SDXL-only prune failed: {e}")
return "\n".join(logs)
# =============================================================================
# STORAGE-EFFICIENT LINEART PIPELINE
# =============================================================================
class SD15LineArtPipeline:
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.dtype = torch.float16 if torch.cuda.is_available() else torch.float32
self.pipeline = None
self.lineart_detector = None
self.depth_estimator = None
self.dpt_processor = None
self.dpt_model = None
self.last_storage_check = None
self.loaded_lora = None # Track currently loaded LoRA
print(f"🚀 Storage-Efficient LineArt initialized on {self.device}")
def _resolve_lora_weight(self, repo_id: str) -> tuple[str, str]:
"""Download LoRA repo and return (local_dir, weight_name) for a .safetensors file."""
try:
local_dir = snapshot_download(repo_id=repo_id, cache_dir=HF_CACHE_DIR, allow_patterns=["*.safetensors", "*.bin"], repo_type="model", token=HF_TOKEN)
# Preferred filenames in order
preferred = [
"pytorch_lora_weights.safetensors",
"adapter_model.safetensors",
"lora.safetensors"
]
import os
for name in preferred:
candidate = os.path.join(local_dir, name)
if os.path.exists(candidate):
return local_dir, name
# Otherwise pick first .safetensors
for root, _, files in os.walk(local_dir):
for f in files:
if f.endswith(".safetensors"):
# Use path relative to local_dir for weight_name
rel = os.path.relpath(os.path.join(root, f), start=local_dir)
return local_dir, rel
# As a last resort, try adapter_model.bin
for root, _, files in os.walk(local_dir):
for f in files:
if f.endswith("adapter_model.bin"):
rel = os.path.relpath(os.path.join(root, f), start=local_dir)
return local_dir, rel
except Exception as e:
print(f"⚠️ snapshot_download failed for {repo_id}: {e}")
return repo_id, None
def check_storage_before_load(self):
"""Check storage before loading models"""
info = get_storage_info()
print(f"📊 Storage check: {info['free_gb']:.1f}GB free, {info['usage_percent']:.1f}% used")
# Clear cache if very low on space
if info['free_gb'] < 1.0:
print("⚠️ Critical storage! Clearing all caches...")
clear_cache_if_needed(min_free_gb=0.5)
self.clear_pipeline()
return False
elif info['free_gb'] < 3.0:
print("⚠️ Low storage! May need to clear cache...")
return True
else:
print("✅ Storage OK for model loading")
return True
def clear_pipeline(self):
"""Clear ALL cached pipelines to free memory"""
print("🗑️ Clearing ALL pipeline caches...")
# Clear main pipeline
if self.pipeline is not None:
print("🗑️ Clearing main pipeline...")
del self.pipeline
self.pipeline = None
# Clear ALL model-specific cached pipelines
attrs_to_clear = []
for attr_name in dir(self):
if attr_name.startswith('pipeline_'):
attrs_to_clear.append(attr_name)
for attr_name in attrs_to_clear:
try:
pipeline_obj = getattr(self, attr_name)
if pipeline_obj is not None:
print(f"🗑️ Clearing cached {attr_name}...")
del pipeline_obj
delattr(self, attr_name)
except (AttributeError, TypeError):
pass
# Clear detectors
if self.lineart_detector is not None:
print("🗑️ Clearing LineArt detector...")
del self.lineart_detector
self.lineart_detector = None
if self.depth_estimator is not None:
print("🗑️ Clearing depth estimator...")
del self.depth_estimator
self.depth_estimator = None
if hasattr(self, 'dpt_processor') and self.dpt_processor is not None:
del self.dpt_processor
self.dpt_processor = None
if hasattr(self, 'dpt_model') and self.dpt_model is not None:
print("🗑️ Clearing DPT model...")
del self.dpt_model
self.dpt_model = None
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
def load_lineart_detector(self):
"""Load LineArt detector with storage check and fallback"""
if self.lineart_detector is not None:
return self.lineart_detector
if LineartDetector is None:
print("❌ LineartDetector not available")
return None
# Check storage before loading
if not self.check_storage_before_load():
return None
# Try multiple approaches for LineArt detector
try:
print("🎨 Loading LineArt Detector...")
# Try 1: Default from_pretrained
try:
self.lineart_detector = LineartDetector.from_pretrained(
"lllyasviel/Annotators",
cache_dir=HF_CACHE_DIR
)
print("✅ LineArt Detector loaded (method 1)!")
return self.lineart_detector
except Exception as e1:
print(f"⚠️ Method 1 failed: {e1}")
# Try 2: Direct instantiation
try:
self.lineart_detector = LineartDetector()
print("✅ LineArt Detector loaded (method 2)!")
return self.lineart_detector
except Exception as e2:
print(f"⚠️ Method 2 failed: {e2}")
# Try 3: Alternative approach without cache_dir
try:
self.lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
print("✅ LineArt Detector loaded (method 3)!")
return self.lineart_detector
except Exception as e3:
print(f"⚠️ Method 3 failed: {e3}")
print("❌ All LineArt Detector loading methods failed")
return None
except Exception as e:
print(f"❌ LineArt Detector loading failed: {str(e)}")
return None
def load_pidinet_detector(self):
"""Load soft-edge PidiNet detector (preferred for SDXL/Union)."""
try:
if PidiNetDetector is None:
print("⚠️ PidiNetDetector not available; falling back to LineArt detector")
return self.load_lineart_detector()
print("🔮 Loading PidiNet soft-edge detector...")
pidinet = PidiNetDetector.from_pretrained("lllyasviel/Annotators", cache_dir=HF_CACHE_DIR)
print("✅ PidiNet detector loaded")
return pidinet
except Exception as e:
print(f"⚠️ PidiNet load failed: {e}; falling back to LineArt detector")
return self.load_lineart_detector()
def load_anime_lineart_detector(self):
"""Load anime-focused LineArt detector."""
try:
if LineartAnimeDetector is None:
print("⚠️ LineartAnimeDetector not available; falling back to LineArt detector")
return self.load_lineart_detector()
print("🎌 Loading Anime LineArt detector...")
anime_det = LineartAnimeDetector.from_pretrained("lllyasviel/Annotators", cache_dir=HF_CACHE_DIR)
print("✅ Anime LineArt detector loaded")
return anime_det
except Exception as e:
print(f"⚠️ Anime detector load failed: {e}; falling back to LineArt detector")
return self.load_lineart_detector()
def load_depth_preprocessor(self):
"""Load depth preprocessor with fallback options"""
if self.depth_estimator is not None:
return self.depth_estimator
# Check storage before loading
if not self.check_storage_before_load():
return None
try:
print("🏔️ Loading Depth Estimator...")
# Try ControlNet-aux MidasDetector first
if DepthEstimator is not None:
try:
self.depth_estimator = DepthEstimator.from_pretrained(
"lllyasviel/Annotators",
cache_dir=HF_CACHE_DIR
)
print("✅ MidasDetector (depth) loaded!")
return self.depth_estimator
except Exception as e1:
print(f"⚠️ MidasDetector failed: {e1}")
# Try direct instantiation if from_pretrained fails
try:
self.depth_estimator = DepthEstimator()
print("✅ MidasDetector (depth) loaded with direct instantiation!")
return self.depth_estimator
except Exception as e2:
print(f"⚠️ MidasDetector direct instantiation failed: {e2}")
# Fallback to DPT model
if DPTImageProcessor is not None and DPTForDepthEstimation is not None:
try:
self.dpt_processor = DPTImageProcessor.from_pretrained(
"Intel/dpt-hybrid-midas",
cache_dir=HF_CACHE_DIR
)
self.dpt_model = DPTForDepthEstimation.from_pretrained(
"Intel/dpt-hybrid-midas",
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR
).to(self.device)
print("✅ DPT depth model loaded as fallback!")
return "dpt_fallback"
except Exception as e2:
print(f"⚠️ DPT fallback failed: {e2}")
print("❌ All depth preprocessor loading methods failed")
return None
except Exception as e:
print(f"❌ Depth preprocessor loading failed: {str(e)}")
return None
def process_depth_map(self, input_image):
"""Process image to depth map"""
try:
if self.depth_estimator is not None:
# Use ControlNet-aux MidasDetector
depth_image = self.depth_estimator(input_image)
print("✅ Depth processed with MidasDetector")
return depth_image
elif self.dpt_model is not None and self.dpt_processor is not None:
# Use DPT fallback
inputs = self.dpt_processor(images=input_image, return_tensors="pt").to(self.device)
with torch.no_grad():
outputs = self.dpt_model(**inputs)
predicted_depth = outputs.predicted_depth
# Convert to PIL Image
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
size=input_image.size[::-1],
mode="bicubic",
align_corners=False,
)
output = prediction.squeeze().cpu().numpy()
formatted = (output * 255 / np.max(output)).astype("uint8")
depth_image = Image.fromarray(formatted).convert("RGB")
print("✅ Depth processed with DPT fallback")
return depth_image
else:
print("❌ No depth preprocessor available")
return None
except Exception as e:
print(f"❌ Depth processing failed: {str(e)}")
return None
def load_lora(self, lora_key="none", strength: float = 0.7):
"""Load and apply LoRA model for coloring book style"""
try:
if not self.pipeline:
print("⚠️ Pipeline not loaded, skipping LoRA")
return False
print(f"🔍 Pipeline LoRA capabilities: has set_adapters={hasattr(self.pipeline,'set_adapters')}, has fuse_lora={hasattr(self.pipeline,'fuse_lora')}, has unload_lora_weights={hasattr(self.pipeline,'unload_lora_weights')}")
lora_config = COLORING_BOOK_LORAS.get(lora_key)
if not lora_config or lora_config["path"] is None:
# Unload LoRA if "none" selected
if self.loaded_lora:
try:
if hasattr(self.pipeline, "unload_lora_weights"):
self.pipeline.unload_lora_weights()
elif hasattr(self.pipeline, "unfuse_lora"):
self.pipeline.unfuse_lora()
print("✅ LoRA unloaded")
except Exception as e:
print(f"⚠️ LoRA unload failed: {e}")
finally:
self.loaded_lora = None
return True
# Skip if same LoRA and same strength already loaded
if self.loaded_lora == f"{lora_key}@{strength:.2f}":
print(f"✅ LoRA '{lora_config['name']}' already loaded at {strength:.2f}")
return True
# Unload previous LoRA if any
if self.loaded_lora:
try:
if hasattr(self.pipeline, "unload_lora_weights"):
self.pipeline.unload_lora_weights()
elif hasattr(self.pipeline, "unfuse_lora"):
self.pipeline.unfuse_lora()
except Exception as e:
print(f"⚠️ Previous LoRA unload failed: {e}")
print(f"🎨 Loading LoRA: {lora_config['name']} @ {strength:.2f}")
adapter_name = "cb_adapter"
local_dir, weight_name = self._resolve_lora_weight(lora_config["path"])
print(f"📦 Resolved LoRA repo: dir={local_dir}, weight={weight_name}")
# 1) Try modern adapters API with adapter_name
try:
if weight_name:
self.pipeline.load_lora_weights(
local_dir,
weight_name=weight_name,
adapter_name=adapter_name
)
else:
self.pipeline.load_lora_weights(
local_dir,
adapter_name=adapter_name
)
if hasattr(self.pipeline, "set_adapters"):
try:
self.pipeline.set_adapters([adapter_name], adapter_weights=[float(strength)])
print("✅ LoRA set via set_adapters")
except Exception as se:
print(f"⚠️ set_adapters(list) failed: {se}; trying scalar signature")
self.pipeline.set_adapters(adapter_name, float(strength))
print("✅ LoRA set via set_adapters (scalar)")
elif hasattr(self.pipeline, "fuse_lora"):
# Older diffusers: manually fuse to apply scale
self.pipeline.fuse_lora(lora_scale=float(strength))
print("✅ LoRA fused into pipeline (no adapters API)")
else:
print("⚠️ Pipeline lacks set_adapters and fuse_lora; LoRA loaded but scale may default to 1.0")
self.loaded_lora = f"{lora_key}@{strength:.2f}"
return True
except Exception as e_adapters:
print(f"⚠️ Adapters API path failed: {e_adapters}")
# 2) Try legacy load with default adapter, then fuse
try:
if weight_name:
self.pipeline.load_lora_weights(local_dir, weight_name=weight_name)
else:
self.pipeline.load_lora_weights(local_dir)
if hasattr(self.pipeline, "set_adapters"):
try:
self.pipeline.set_adapters(["default"], adapter_weights=[float(strength)])
print("✅ LoRA set via default adapter")
except Exception as se2:
print(f"⚠️ set_adapters(default) failed: {se2}")
raise se2
elif hasattr(self.pipeline, "fuse_lora"):
self.pipeline.fuse_lora(lora_scale=float(strength))
else:
print("⚠️ Pipeline lacks set_adapters/fuse_lora in legacy path")
self.loaded_lora = f"{lora_key}@{strength:.2f}"
return True
except Exception as e_legacy:
print(f"❌ Legacy LoRA load failed: {e_legacy}")
self.loaded_lora = None
return False
except Exception as e:
print(f"❌ LoRA loading failed: {e}")
self.loaded_lora = None
return False
def apply_pure_bw_post_processing(self, image):
"""Apply AGGRESSIVE edge detection and thresholding for pure coloring book style"""
try:
import cv2
# Convert PIL to numpy
img_array = np.array(image.convert('RGB'))
# Convert to grayscale
gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
# Apply stronger bilateral filter to preserve edges while removing noise
filtered = cv2.bilateralFilter(gray, 9, 80, 80)
# Apply adaptive threshold for better line detection
adaptive_thresh = cv2.adaptiveThreshold(
filtered, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 11, 10
)
# Apply more aggressive Canny edge detection
edges = cv2.Canny(filtered, 30, 100)
# Use larger kernel for morphological operations to strengthen lines
kernel = np.ones((3, 3), np.uint8)
# Close gaps in lines
edges_closed = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
# Dilate to thicken lines slightly
edges_thick = cv2.dilate(edges_closed, kernel, iterations=1)
# Combine adaptive threshold with edge detection
combined = cv2.bitwise_or(adaptive_thresh, edges_thick)
# Apply final binary threshold to ensure pure black and white
_, final_binary = cv2.threshold(combined, 127, 255, cv2.THRESH_BINARY)
# Invert so lines are black on white background
inverted = cv2.bitwise_not(final_binary)
# Remove small noise spots
kernel_small = np.ones((2, 2), np.uint8)
cleaned = cv2.morphologyEx(inverted, cv2.MORPH_OPEN, kernel_small)
# Convert back to PIL
result_image = Image.fromarray(cleaned).convert('RGB')
print("✅ AGGRESSIVE coloring book post-processing applied")
return result_image
except Exception as e:
print(f"⚠️ Post-processing failed: {e}, returning original")
return image
def load_pipeline(self, lineart_model_key="standard", pipeline_type: str = "sd15"):
"""Load pipeline with storage management and better error handling"""
# Create a unique pipeline key that includes the model selection
pipeline_key = f"pipeline_{pipeline_type}_{lineart_model_key}"
# Check if we have this specific pipeline cached
if hasattr(self, pipeline_key) and getattr(self, pipeline_key) is not None:
return getattr(self, pipeline_key)
# Clear cache if needed BEFORE loading
if not clear_cache_if_needed(min_free_gb=6.0):
print("❌ Insufficient storage even after cleanup")
return None
# Check storage before loading
if not self.check_storage_before_load():
return None
try:
if pipeline_type == "sdxl":
model_config = LINEART_MODELS_SDXL.get(lineart_model_key, LINEART_MODELS_SDXL["union"])
model_type = model_config["type"]
model_id = model_config["id"]
print(f"🎨 Loading SDXL {model_type.upper()} model: {model_id}")
print(f"📊 Storage before load: {get_storage_info()['free_gb']:.1f}GB free")
# Load Union ControlNet (SDXL)
try:
controlnet = ControlNetModel.from_pretrained(
model_id,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR,
use_safetensors=True
)
print("✅ SDXL ControlNet loaded successfully")
except Exception as controlnet_error:
print(f"❌ SDXL ControlNet loading failed: {controlnet_error}")
print("🔄 Trying fallback loading (no safetensors)...")
controlnet = ControlNetModel.from_pretrained(
model_id,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR
)
print("✅ SDXL ControlNet loaded with fallback method")
# Load SDXL base with ControlNet
print(f"⏳ Loading SDXL base model: {SDXL_MODEL_ID}...")
try:
pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
SDXL_MODEL_ID,
controlnet=controlnet,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR,
use_safetensors=True
)
print("✅ SDXL Pipeline loaded successfully")
except Exception as pipeline_error:
print(f"❌ SDXL Pipeline loading failed: {pipeline_error}")
pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
SDXL_MODEL_ID,
controlnet=controlnet,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR
)
print("✅ SDXL Pipeline loaded with fallback")
else:
model_config = LINEART_MODELS.get(lineart_model_key, LINEART_MODELS["standard"])
model_type = model_config["type"]
model_id = model_config["id"]
print(f"🎨 Loading {model_type.upper()} model: {model_id}")
print(f"📊 Storage before load: {get_storage_info()['free_gb']:.1f}GB free")
# All models are ControlNet now (simpler!)
if model_type == "controlnet":
print(f"⏳ Loading ControlNet from {model_id}...")
# Try to load ControlNet with better error handling
try:
controlnet = ControlNetModel.from_pretrained(
model_id,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR,
use_safetensors=True
)
print("✅ ControlNet loaded successfully")
except Exception as controlnet_error:
print(f"❌ ControlNet loading failed: {controlnet_error}")
# Try fallback: simpler loading without safetensors
print("🔄 Trying fallback loading...")
controlnet = ControlNetModel.from_pretrained(
model_id,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR
)
print("✅ ControlNet loaded with fallback method")
print(f"⏳ Loading SD 1.5 base model: {SD15_MODEL_ID}...")
# Create SD 1.5 ControlNet pipeline (much simpler!)
try:
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
SD15_MODEL_ID,
controlnet=controlnet,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR,
use_safetensors=True
)
print("✅ SD 1.5 Pipeline loaded successfully")
except Exception as pipeline_error:
print(f"❌ Pipeline loading failed: {pipeline_error}")
# Try without safetensors
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
SD15_MODEL_ID,
controlnet=controlnet,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR
)
print("✅ SD 1.5 Pipeline loaded with fallback")
else:
raise ValueError(f"Unknown model type: {model_type}")
# Check storage after model loading
info = get_storage_info()
print(f"📊 After model load: {info['free_gb']:.1f}GB free")
if info['free_gb'] < 1.0:
print("❌ Not enough space for full pipeline!")
return None
# Basic optimizations
pipeline.enable_model_cpu_offload()
# Set active pipeline reference for LoRA loading
self.pipeline = pipeline
# Cache this specific pipeline
setattr(self, pipeline_key, pipeline)
# Final storage check
info = get_storage_info()
print(f"📊 After pipeline load: {info['free_gb']:.1f}GB free")
print(f"✅ {('SDXL' if pipeline_type=='sdxl' else model_type.upper())} pipeline loaded successfully for {lineart_model_key}!")
return pipeline
except Exception as e:
print(f"❌ Pipeline loading failed: {str(e)}")
# Try to clear cache and return None
clear_cache_if_needed(min_free_gb=0.5)
return None
def generate_with_lineart(self, input_image, prompt="coloring book page, black and white line art",
num_steps=20, guidance_scale=7.5, controlnet_scale=1.2,
style_preset="clean", use_negative_prompt=True, output_resolution=768,
use_depth_preprocessor=False, lineart_model="standard", apply_post_processing=False,
lora_model="none", pipeline_type: str = "sd15", lora_strength: float = 0.7,
monochrome_lock: bool = True):
"""Generate with LineArt and advanced controls for coloring book optimization"""
try:
# Select detector based on chosen model
if lineart_model == "soft_edge":
detector = self.load_pidinet_detector()
elif lineart_model == "anime":
detector = self.load_anime_lineart_detector()
elif lineart_model == "canny":
detector = None # will fall back to Canny below
else:
detector = self.load_lineart_detector()
# Process input image to LineArt
input_image = Image.fromarray(input_image).convert("RGB")
if detector is None:
print("⚠️ LineArt Detector failed, using simple edge detection fallback...")
# Simple fallback: convert to grayscale and apply basic edge detection
import cv2
gray = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2GRAY)
edges = cv2.Canny(gray, 50, 150)
# Convert back to 3-channel for consistency
lineart_image = Image.fromarray(cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB))
print("✅ Fallback edge detection complete")
else:
print("🎨 Processing image with LineArt detector...")
lineart_image = detector(input_image)
print("✅ LineArt processing complete")
# Load pipeline with selected LineArt model and pipeline type
pipeline = self.load_pipeline(
lineart_model_key=(lineart_model if pipeline_type == "sd15" else "union"),
pipeline_type=pipeline_type
)
if pipeline is None:
return None, lineart_image, None, "❌ Pipeline failed to load (storage full?)"
# Load LoRA model if specified
lora_status = ""
if lora_model != "none":
print(f"🎨 Loading LoRA: {lora_model}")
lora_success = self.load_lora(lora_model, strength=lora_strength)
if lora_success:
lora_status = f" | LoRA: {COLORING_BOOK_LORAS[lora_model]['name']} @ {lora_strength:.2f}"
else:
print(f"⚠️ LoRA loading failed, continuing without LoRA")
else:
# Explicitly unload any prior LoRA if user selected None
self.load_lora("none")
# Check storage before generation
info = get_storage_info()
if info['free_gb'] < 0.5:
return None, lineart_image, None, f"❌ Insufficient storage for generation: {info['free_gb']:.1f}GB free"
# ENHANCED prompts for PURE coloring book style
style_prompts = {
"ultra_clean": "simple line drawing, black lines on white background, coloring book style, vector art, minimalist line art, clean outlines only, no shading, no gradients, no textures, flat design, monochrome, black and white",
"clean": "coloring book line art, simple black outlines, clean line drawing, minimal details, bold lines, no shading, white background, easy to color, monochrome",
"detailed": "detailed line art, coloring book page, intricate line drawing, fine black lines, no shading, complex outlines, black and white only, monochrome",
"bold": "thick black lines, bold outlines, heavy line weight, strong contrast, coloring book style, simple shapes, no shading, monochrome",
"simple": "basic line drawing, simple outlines, beginner coloring book, large shapes, thick lines, minimal details, easy to color, monochrome"
}
if monochrome_lock:
style_prompts = {k: v + ", pure black-and-white, grayscale only, no colored areas, no fills" for k, v in style_prompts.items()}
# Get style-specific prompt additions
style_addition = style_prompts.get(style_preset, style_prompts["clean"])
# Add anime-specific styling if anime model selected
if lineart_model == "anime":
anime_enhancement = "anime line art, manga style lineart, simple anime drawing, anime coloring book, clean anime outlines, japanese cartoon style, anime character design, line art only"
style_addition = f"{anime_enhancement}, {style_addition}"
# Add LoRA trigger words if LoRA is loaded
lora_trigger = ""
if lora_model != "none" and lora_model in COLORING_BOOK_LORAS:
lora_config = COLORING_BOOK_LORAS[lora_model]
if lora_config["trigger"]:
lora_trigger = f"{lora_config['trigger']}, "
print(f"🎯 Using LoRA trigger: {lora_config['trigger']}")
final_prompt = f"{lora_trigger}{style_addition}, {prompt}" if prompt.strip() else f"{lora_trigger}{style_addition}"
# AGGRESSIVE negative prompts for PURE coloring book style
negative_prompts = {
"ultra_clean": "color, colors, colored, colorized, saturation, vivid, pastel, fill, filled areas, shading, shadows, gradients, realistic, photographic, photography, photorealistic, 3d render, painting, watercolor, oil painting, detailed textures, complex lighting, depth, perspective, artistic style, sketch, pencil drawing, charcoal, hatching, crosshatching, tones, gray, grey, backgrounds, textures, noise, blur, soft edges, detailed faces, realistic skin, realistic hair, complex details, fine details, intricate textures",
"clean": "color, colors, colored, colorized, saturation, pastel, fill, shading, shadows, gradients, realistic, photographic, photorealistic, complex details, artistic style, sketch, pencil drawing, tones, filled areas, textures, detailed faces, realistic features",
"detailed": "color, colors, colored, colorized, heavy shading, shadows, realistic, photographic, photorealistic, gradients, filled areas, sketch style, textures, soft lines",
"bold": "color, colors, colored, colorized, shading, shadows, realistic, photographic, photorealistic, complex details, gradients, thin lines, weak lines, soft edges, blur",
"simple": "color, colors, colored, colorized, shading, shadows, realistic, photographic, photorealistic, intricate patterns, fine details, textures, detailed features"
}
# SDXL needs extra reinforcement against colorization
if pipeline_type == "sdxl" and use_negative_prompt:
extra = ", skin tone, skin color, skin, color cast, color bleed, chroma, hue, saturation, tint, tone"
if monochrome_lock:
extra += ", color, colored, colorized, pastel fill, marker, watercolor, paint, pigment"
negative_prompts = {k: v + extra for k, v in negative_prompts.items()}
final_negative = negative_prompts.get(style_preset, negative_prompts["clean"]) if use_negative_prompt else ""
# Add anime-specific negative prompts if anime model selected
if lineart_model == "anime" and use_negative_prompt:
anime_negative = "western style, realistic faces, photorealistic faces, 3D style, realistic proportions, detailed facial features, complex shading, realistic lighting, western art style"
final_negative = f"{final_negative}, {anime_negative}" if final_negative else anime_negative
print(f"🎨 Generating with style: {style_preset}")
print(f"📝 Prompt: {final_prompt[:100]}...")
print(f"🚫 Negative: {final_negative[:50]}...")
# Calculate output dimensions
aspect_ratio = input_image.width / input_image.height
if aspect_ratio > 1: # Landscape
output_width = min(output_resolution, 1536) # Max 1536 for memory
output_height = int(output_width / aspect_ratio)
else: # Portrait or square
output_height = min(output_resolution, 1536)
output_width = int(output_height * aspect_ratio)
# Ensure dimensions are multiples of 8 (required by SDXL)
output_width = (output_width // 8) * 8
output_height = (output_height // 8) * 8
print(f"🖼️ Output resolution: {output_width}x{output_height}")
# All models use ControlNet pipeline now (simpler!)
result = pipeline(
prompt=final_prompt,
negative_prompt=final_negative,
image=lineart_image,
num_inference_steps=num_steps,
guidance_scale=guidance_scale,
controlnet_conditioning_scale=controlnet_scale,
width=output_width,
height=output_height
)
# Get the generated image
generated_image = result.images[0]
# Apply post-processing if enabled
if apply_post_processing:
generated_image = self.apply_pure_bw_post_processing(generated_image)
# Process depth if enabled
depth_image = None
if use_depth_preprocessor:
depth_estimator = self.load_depth_preprocessor()
if depth_estimator:
depth_image = self.process_depth_map(input_image)
if depth_image:
print("✅ Depth map generated for display")
post_processing_text = " + Post-Processing" if apply_post_processing else ""
print("✅ Generation complete")
return generated_image, lineart_image, depth_image, f"✅ Generated with {style_preset} style, {num_steps} steps{lora_status}{post_processing_text}!"
except Exception as e:
error_msg = f"❌ Generation failed: {str(e)}"
print(error_msg)
# If it's a storage error, try to clear cache
if "out of memory" in str(e).lower() or "space" in str(e).lower():
clear_cache_if_needed(min_free_gb=0.5)
self.clear_pipeline()
error_msg += " (Cleared cache due to storage issue)"
return None, None, None, error_msg
def prefetch_loras(self) -> str:
"""Download all configured LoRAs into persistent storage and validate files."""
logs = ["📦 Prefetching LoRAs to persistent storage..."]
if HF_TOKEN:
logs.append("🔑 Using HF_TOKEN for authenticated downloads")
else:
logs.append("ℹ️ No HF_TOKEN set; private repos will be skipped")
import os
total_ok = 0
for key, cfg in COLORING_BOOK_LORAS.items():
if not cfg.get("path"):
logs.append(f"⏭️ Skipping {cfg['name']}: disabled or no repo path")
continue
try:
local_dir = snapshot_download(repo_id=cfg["path"], cache_dir=HF_CACHE_DIR, allow_patterns=["*.safetensors", "*.bin"], repo_type="model", token=HF_TOKEN)
found = []
for root, _, files in os.walk(local_dir):
for f in files:
if f.endswith(".safetensors") or f.endswith("adapter_model.bin"):
p = os.path.join(root, f)
size_mb = os.path.getsize(p) / (1024*1024)
found.append((p, size_mb))
if found:
lines = [f"✅ {cfg['name']}: {len(found)} files"] + [f" - {os.path.relpath(p, HF_CACHE_DIR)} ({size_mb:.1f}MB)" for p, size_mb in found]
logs.extend(lines)
total_ok += 1
else:
logs.append(f"⚠️ {cfg['name']}: no weight files found after download")
except Exception as e:
logs.append(f"❌ {cfg['name']}: download failed - {e}")
logs.append(f"📊 LoRA prefetch complete: {total_ok} repos ready. Cache dir: {HF_CACHE_DIR}")
return "\n".join(logs)
# =============================================================================
# GRADIO INTERFACE WITH STORAGE MONITORING
# =============================================================================
def create_storage_efficient_interface():
"""Create storage-efficient interface with monitoring"""
# Initialize pipeline
pipeline_manager = SD15LineArtPipeline()
def generate_with_monitoring(input_image, custom_prompt, style_preset, num_steps,
guidance_scale, controlnet_scale, use_negative_prompt,
output_resolution, use_depth_preprocessor, lineart_model, apply_post_processing, lora_model, pipeline_type, lora_strength, monochrome_lock):
if input_image is None:
return None, None, None, "Please upload an image", format_storage_status()
# Use custom prompt or default
prompt = custom_prompt.strip() if custom_prompt.strip() else ""
# Generate with all user controls
result_image, lineart_preview, depth_preview, status = pipeline_manager.generate_with_lineart(
input_image,
prompt,
num_steps=num_steps,
guidance_scale=guidance_scale,
controlnet_scale=controlnet_scale,
style_preset=style_preset,
use_negative_prompt=use_negative_prompt,
output_resolution=output_resolution,
use_depth_preprocessor=use_depth_preprocessor,
lineart_model=lineart_model,
apply_post_processing=apply_post_processing,
lora_model=lora_model,
pipeline_type=pipeline_type,
lora_strength=lora_strength,
monochrome_lock=monochrome_lock
)
# Update storage status
storage_status = format_storage_status()
return result_image, lineart_preview, depth_preview, status, storage_status
def manual_clear_cache():
pipeline_manager.clear_pipeline()
cleared = clear_cache_if_needed(min_free_gb=0.0) # Force clear
status = "✅ Cache cleared!" if cleared else "✅ Cache clear attempted"
storage_status = format_storage_status()
return status, storage_status
def prefetch_loras_action():
msg = pipeline_manager.prefetch_loras()
return msg, format_storage_status()
def prune_sdxl_only_action(keep_depth: bool):
# Keep LoRA repos that are enabled
lora_repo_ids = [cfg["path"] for cfg in COLORING_BOOK_LORAS.values() if cfg.get("path")]
log = prune_cache_to_sdxl_only(keep_depth=keep_depth, keep_lora_repo_ids=lora_repo_ids)
return log, format_storage_status()
def emergency_full_cleanup():
"""Emergency full cleanup - removes everything for SD 1.5 migration"""
try:
print("🚨 EMERGENCY CLEANUP: Full cache wipe for SD 1.5...")
# Clear all pipelines first
pipeline_manager.clear_pipeline()
# Clear HF cache completely
if os.path.exists(HF_CACHE_DIR):
shutil.rmtree(HF_CACHE_DIR)
os.makedirs(HF_CACHE_DIR, exist_ok=True)
print("🗑️ HF cache wiped")
# Clear any models in /data (old SDXL models)
if os.path.exists("/data"):
for item in os.listdir("/data"):
if any(pattern in item for pattern in [
"models--stabilityai", "models--diffusers", "models--madebyollin",
"models--Intel", "hub", "transformers", "diffusers"
]):
try:
item_path = os.path.join("/data", item)
if os.path.isdir(item_path):
shutil.rmtree(item_path)
else:
os.remove(item_path)
print(f"🗑️ Removed: {item}")
except Exception as e:
print(f"⚠️ Could not remove {item}: {e}")
# Force garbage collection
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
status = "🧹 EMERGENCY CLEANUP COMPLETE! Ready for SD 1.5 models."
storage_status = format_storage_status()
print("✅ Emergency cleanup complete!")
return status, storage_status
except Exception as e:
error_msg = f"❌ Emergency cleanup failed: {e}"
print(error_msg)
return error_msg, format_storage_status()
def on_pipeline_change(pipeline_choice):
if pipeline_choice == "sdxl":
return (
gr.update(maximum=1536, value=1024, info="SDXL optimal: 1024px; up to 1536px"),
gr.update(value=6.5),
gr.update(value=1.8),
gr.update(value=0.7)
)
else:
return (
gr.update(maximum=768, value=512, info="SD 1.5 optimal: 512px"),
gr.update(value=8.0),
gr.update(value=1.2),
gr.update(value=0.8)
)
with gr.Blocks(title="ColorCraft SDXL - Advanced LineArt Controls") as demo:
gr.HTML("""
🎨 ColorCraft SDXL - Advanced LineArt Controls
Fine-tuned controls for clean black & white coloring books • LineArt preprocessor • Storage monitoring
""")
with gr.Row():
with gr.Column(scale=1):
gr.HTML("📷 Input & Controls
")
input_image = gr.Image(
label="Upload Image",
type="numpy",
height=300
)
# Style Preset - Most Important Control
style_preset = gr.Dropdown(
label="🎨 Style Preset",
choices=["ultra_clean", "clean", "detailed", "bold", "simple"],
value="ultra_clean",
info="Ultra Clean = Pure B&W, Clean = Minimal shading"
)
# LineArt Model Selection
lineart_model = gr.Dropdown(
label="🖌️ LineArt Model",
choices=[
("Standard LineArt (lllyasviel)", "standard"),
("Anime LineArt Specialist (v1.1)", "anime"),
("Soft Edge Detection (Better for Coloring)", "soft_edge"),
("Canny Edge Detection (Harsh)", "canny")
],
value="soft_edge",
info="RECOMMENDED: Soft Edge for cleaner coloring book results"
)
# Pipeline Type Selection (SD 1.5 vs SDXL)
pipeline_type = gr.Dropdown(
label="⚡ Pipeline",
choices=[
("Fast (SD 1.5 - 512px)", "sd15"),
("Professional (SDXL - 1024px)", "sdxl")
],
value="sdxl",
info="Choose SDXL for highest quality (slower, larger models)"
)
# LoRA Model Selection for Coloring Book Enhancement
lora_choices = [("No LoRA", "none")] + [
(cfg["name"], k) for k, cfg in COLORING_BOOK_LORAS.items() if cfg.get("path")
]
lora_model = gr.Dropdown(
label="🎨 Coloring Book LoRA",
choices=lora_choices,
value="coloring_book_redmond",
info="RECOMMENDED: Specialized LoRA models for professional coloring book style"
)
lora_strength = gr.Slider(
label="LoRA Strength",
minimum=0.0,
maximum=1.0,
value=0.7,
step=0.05,
info="Recommended: 0.7 for SDXL, 0.8 for SD 1.5"
)
monochrome_lock = gr.Checkbox(
label="Monochrome Lock (enforce black & white)",
value=True,
info="Adds strong monochrome cues and color blockers for SDXL"
)
custom_prompt = gr.Textbox(
label="Custom Prompt (optional)",
placeholder="Additional prompt (style preset will be added automatically)",
lines=2,
value=""
)
# Fine-tuning Controls
with gr.Accordion("⚙️ Advanced Settings", open=True):
with gr.Row():
num_steps = gr.Slider(
label="Inference Steps",
minimum=10,
maximum=50,
value=25,
step=5,
info="More steps = higher quality, slower"
)
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=3.0,
maximum=15.0,
value=8.0,
step=0.5,
info="OPTIMAL: 7-9 for balanced results (was too aggressive at 12+)"
)
with gr.Row():
controlnet_scale = gr.Slider(
label="ControlNet Strength",
minimum=0.5,
maximum=2.0,
value=1.2,
step=0.1,
info="OPTIMAL: 1.0-1.2 for SD 1.5; SDXL Union works best around 1.7–2.0"
)
use_negative_prompt = gr.Checkbox(
label="Use Negative Prompt",
value=True,
info="Helps remove colors and shading"
)
with gr.Row():
output_resolution = gr.Slider(
label="Output Resolution",
minimum=384,
maximum=1536,
value=1024,
step=8,
info="SDXL optimal: 1024px (up to 1536px). Switch to SD 1.5 for 512px."
)
use_depth_preprocessor = gr.Checkbox(
label="Add Depth Preprocessing",
value=False,
info="Experimental: Add structural depth guidance"
)
with gr.Row():
apply_post_processing = gr.Checkbox(
label="Pure B&W Post-Processing",
value=False,
info="DISABLED: Was making output worse - now using LoRA + optimized prompts instead"
)
generate_btn = gr.Button(
"🎨 Generate LineArt",
variant="primary",
size="lg"
)
with gr.Column(scale=2):
gr.HTML("🎨 Generated Results
")
output_image = gr.Image(
label="Generated Coloring Book",
height=300
)
with gr.Row():
lineart_preview = gr.Image(
label="LineArt Preprocessor Output",
height=200
)
depth_preview = gr.Image(
label="Depth Preprocessor Output",
height=200
)
with gr.Column():
status_output = gr.Textbox(
label="Generation Status",
value="Ready to generate...",
lines=3
)
storage_status = gr.Textbox(
label="Storage Status",
value=format_storage_status(),
lines=2
)
# Storage Management
with gr.Accordion("💾 Storage Management", open=True):
with gr.Row():
clear_cache_btn = gr.Button("🧹 Clear Cache")
clear_all_btn = gr.Button("🧨 EMERGENCY CLEANUP", variant="stop")
prefetch_loras_btn = gr.Button("📦 Prefetch LoRAs to Storage")
refresh_storage_btn = gr.Button("🔄 Refresh Storage Status")
prune_sdxl_btn = gr.Button("🧽 Prune to SDXL-Only Cache")
keep_depth_checkbox = gr.Checkbox(label="Keep Depth Model (Intel/dpt-hybrid-midas)", value=False)
gr.HTML("""
Storage Tips:
• Models require ~8-12GB total
• Generation needs ~2GB free space
• Clear cache if generation fails
• Click refresh to update storage status
""")
# Add style preset guide
with gr.Accordion("📖 Style Guide", open=False):
gr.HTML("""
🎨 Style Presets Explained:
- Ultra Clean: BEST for coloring books - minimal details, pure line art, vector-like
- Clean: Simple outlines, clean lines, easy to color (recommended)
- Detailed: More intricate lines but still coloring book appropriate
- Bold: Thick outlines, high contrast, great for kids
- Simple: Basic outlines, perfect for beginners
⚙️ OPTIMAL Settings for Clean Coloring Books (RESEARCH-BASED):
- 🎨 LoRA: Use "ColoringBook Redmond V2" for professional results
- 🖌️ Model: "Soft Edge Detection" for cleaner lines vs harsh Canny
- Steps: 25-30 (50+ overkill, 25 sufficient with LoRA)
- Guidance: 7-9 balanced (12+ was too aggressive)
- ControlNet: 1.0-1.2 balanced (1.6 was too strong)
- Post-Processing: DISABLED (LoRA + prompts work better)
- Style: "Ultra Clean" preset for minimal details
""")
# Event handlers
generate_btn.click(
generate_with_monitoring,
inputs=[input_image, custom_prompt, style_preset, num_steps,
guidance_scale, controlnet_scale, use_negative_prompt,
output_resolution, use_depth_preprocessor, lineart_model, apply_post_processing, lora_model, pipeline_type, lora_strength, monochrome_lock],
outputs=[output_image, lineart_preview, depth_preview, status_output, storage_status]
)
clear_cache_btn.click(
manual_clear_cache,
outputs=[status_output, storage_status]
)
clear_all_btn.click(
emergency_full_cleanup,
outputs=[status_output, storage_status]
)
prefetch_loras_btn.click(
prefetch_loras_action,
outputs=[status_output, storage_status]
)
prune_sdxl_btn.click(
prune_sdxl_only_action,
inputs=[keep_depth_checkbox],
outputs=[status_output, storage_status]
)
refresh_storage_btn.click(
lambda: format_storage_status(),
outputs=storage_status
)
pipeline_type.change(
on_pipeline_change,
inputs=[pipeline_type],
outputs=[output_resolution, guidance_scale, controlnet_scale, lora_strength]
)
# Note: Auto-update removed due to Gradio compatibility issues
# Storage status updates manually via buttons
return demo
# =============================================================================
# MAIN EXECUTION
# =============================================================================
if __name__ == "__main__":
print("🚀 Starting ColorCraft SD 1.5 - LineArt Specialist Edition")
print(f"💾 Cache directory: {HF_CACHE_DIR}")
# Initial storage check
info = get_storage_info()
print(f"📊 Initial storage: {info['free_gb']:.1f}GB free of {info['total_gb']:.1f}GB total")
# Clear cache immediately on startup for SD 1.5 migration
print("🧹 Startup cache cleanup for SD 1.5 migration...")
clear_cache_if_needed(min_free_gb=8.0)
# Test imports
print("\n🔍 Testing imports...")
try:
print(f"✅ PyTorch: {torch.__version__}")
print(f"✅ Device: {torch.cuda.get_device_name() if torch.cuda.is_available() else 'CPU'}")
import diffusers
print(f"✅ Diffusers: {diffusers.__version__}")
import transformers
print(f"✅ Transformers: {transformers.__version__}")
import huggingface_hub
print(f"✅ HuggingFace Hub: {huggingface_hub.__version__}")
if LineartDetector:
print("✅ ControlNet-Aux: Available")
else:
print("❌ ControlNet-Aux: Not available")
except Exception as e:
print(f"❌ Import test failed: {e}")
# Create and launch interface
demo = create_storage_efficient_interface()
demo.queue(max_size=3) # Smaller queue to save memory
demo.launch(
server_name="0.0.0.0",
server_port=7860,
show_error=True
)