"""
ColorCraft SDXL - LineArt Storage Efficient Edition
==================================================
Storage-efficient LineArt approach:
- On-demand model loading with cache management
- Storage monitoring and cleanup
- LineArt preprocessor output display
- Single model approach to minimize storage usage
"""
import os
import gc
import shutil
import gradio as gr
import torch
import numpy as np
from PIL import Image
# 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 (
StableDiffusionXLControlNetPipeline,
StableDiffusionXLAdapterPipeline,
ControlNetModel,
T2IAdapter,
AutoencoderKL,
DDIMScheduler
)
print("✅ Diffusers imported successfully")
except ImportError as e:
print(f"❌ Diffusers import error: {e}")
# Try ControlNet preprocessor imports
try:
from controlnet_aux import LineartDetector
print("✅ LineartDetector imported successfully")
except ImportError as e:
print(f"❌ LineartDetector import error: {e}")
LineartDetector = 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
# =============================================================================
SDXL_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
VAE_MODEL_ID = "madebyollin/sdxl-vae-fp16-fix"
# LineArt model options (verified SDXL compatible)
LINEART_MODELS = {
"standard": {
"type": "controlnet",
"id": "ShermanG/ControlNet-Standard-Lineart-for-SDXL"
},
"anime": {
"type": "controlnet",
"id": "r3gm/controlnet-lineart-anime-sdxl-fp16"
},
"tencent": {
"type": "t2i_adapter",
"id": "TencentARC/t2i-adapter-lineart-sdxl-1.0"
}
}
DEPTH_MODEL_ID = "diffusers/controlnet-depth-sdxl-1.0"
# 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)
# =============================================================================
# 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=2.0):
"""Clear cache if storage is low"""
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
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}")
# Force garbage collection
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
return True
return False
# =============================================================================
# STORAGE-EFFICIENT LINEART PIPELINE
# =============================================================================
class StorageEfficientLineArtPipeline:
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
print(f"🚀 Storage-Efficient LineArt initialized on {self.device}")
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 current pipeline to free memory"""
if self.pipeline is not None:
print("🗑️ Clearing pipeline to free storage...")
del self.pipeline
self.pipeline = None
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 self.dpt_model is not None:
print("🗑️ Clearing DPT model...")
del self.dpt_model
self.dpt_model = None
if self.dpt_processor is not None:
del self.dpt_processor
self.dpt_processor = 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_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 apply_pure_bw_post_processing(self, image):
"""Apply edge detection and thresholding for pure black and white"""
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 Gaussian blur to reduce noise
blurred = cv2.GaussianBlur(gray, (3, 3), 0)
# Apply Canny edge detection
edges = cv2.Canny(blurred, 50, 150)
# Apply morphological operations to clean up lines
kernel = np.ones((2, 2), np.uint8)
edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
# Apply binary thresholding for pure black and white
_, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
# Combine edges with binary threshold
combined = cv2.bitwise_or(edges, binary)
# Invert so lines are black on white background
inverted = cv2.bitwise_not(combined)
# Convert back to PIL
result_image = Image.fromarray(inverted).convert('RGB')
print("✅ Pure B&W 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"):
"""Load pipeline with storage management - supports both ControlNet and T2I-Adapter"""
# Create a unique pipeline key that includes the model selection
pipeline_key = f"pipeline_{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)
# Check storage before loading
if not self.check_storage_before_load():
return None
try:
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}")
# Load VAE first
vae = AutoencoderKL.from_pretrained(
VAE_MODEL_ID,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR
)
if model_type == "controlnet":
# Load ControlNet model
controlnet = ControlNetModel.from_pretrained(
model_id,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR,
variant="fp16" if self.dtype == torch.float16 else None
)
# Create ControlNet pipeline
pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
SDXL_MODEL_ID,
controlnet=controlnet,
vae=vae,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR,
use_safetensors=True,
variant="fp16" if self.dtype == torch.float16 else None
)
elif model_type == "t2i_adapter":
# Load T2I-Adapter model
adapter = T2IAdapter.from_pretrained(
model_id,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR,
variant="fp16" if self.dtype == torch.float16 else None
)
# Create T2I-Adapter pipeline
pipeline = StableDiffusionXLAdapterPipeline.from_pretrained(
SDXL_MODEL_ID,
adapter=adapter,
vae=vae,
torch_dtype=self.dtype,
cache_dir=HF_CACHE_DIR,
use_safetensors=True,
variant="fp16" if self.dtype == torch.float16 else None
)
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()
# 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"✅ {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):
"""Generate with LineArt and advanced controls for coloring book optimization"""
try:
# Load detector first
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
pipeline = self.load_pipeline(lineart_model)
if pipeline is None:
return None, lineart_image, None, "❌ Pipeline failed to load (storage full?)"
# 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"
# Build optimized prompts based on style preset
style_prompts = {
"ultra_clean": "pure black and white line art, coloring book page, simple outlines, minimal details, vector art style, clean lines, no shading, no colors, stark contrast",
"clean": "black and white line art, coloring book page, clear outlines, simple design, minimal shading, clean vector style",
"detailed": "detailed black and white line art, coloring book page, intricate outlines, fine details, precise lines",
"bold": "bold black and white line art, coloring book page, thick outlines, strong lines, high contrast",
"simple": "simple black and white line art, coloring book page, basic outlines, easy to color"
}
# Get style-specific prompt additions
style_addition = style_prompts.get(style_preset, style_prompts["clean"])
final_prompt = f"{style_addition}, {prompt}" if prompt.strip() else style_addition
# Optimized negative prompts for coloring book style
negative_prompts = {
"ultra_clean": "color, colors, colored, shading, shadows, gradients, realistic, photographic, photography, 3d render, painting, watercolor, oil painting, detailed textures, complex lighting, depth, perspective, artistic style, sketch",
"clean": "color, colors, shading, gradients, realistic, photographic, complex details, artistic style",
"detailed": "color, colors, heavy shading, realistic, photographic",
"bold": "color, colors, fine details, realistic, photographic",
"simple": "color, colors, complex details, shading, realistic, photographic"
}
final_negative = negative_prompts.get(style_preset, negative_prompts["clean"]) if use_negative_prompt else ""
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}")
# Get model type to determine parameter names
model_config = LINEART_MODELS.get(lineart_model, LINEART_MODELS["standard"])
model_type = model_config["type"]
# Generate with type-specific parameters
if model_type == "controlnet":
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
)
elif model_type == "t2i_adapter":
result = pipeline(
prompt=final_prompt,
negative_prompt=final_negative,
image=lineart_image,
num_inference_steps=num_steps,
guidance_scale=guidance_scale,
adapter_conditioning_scale=controlnet_scale, # Different parameter name for T2I-Adapter
width=output_width,
height=output_height
)
else:
raise ValueError(f"Unknown model type: {model_type}")
# 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{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
# =============================================================================
# GRADIO INTERFACE WITH STORAGE MONITORING
# =============================================================================
def create_storage_efficient_interface():
"""Create storage-efficient interface with monitoring"""
# Initialize pipeline
pipeline_manager = StorageEfficientLineArtPipeline()
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):
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
)
# 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
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 ControlNet", "standard"),
("Anime ControlNet (r3gm)", "anime"),
("TencentARC T2I-Adapter", "tencent")
],
value="standard",
info="Test different LineArt approaches: ControlNet vs T2I-Adapter architectures"
)
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=9.0,
step=0.5,
info="Higher = follows prompt more strictly"
)
with gr.Row():
controlnet_scale = gr.Slider(
label="ControlNet Strength",
minimum=0.5,
maximum=2.0,
value=1.4,
step=0.1,
info="Higher = follows input lines more closely"
)
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=512,
maximum=1536,
value=1024,
step=64,
info="Higher = better quality, slower generation"
)
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="Apply edge detection + thresholding for pure black/white"
)
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 & Free Storage")
refresh_storage_btn = gr.Button("🔄 Refresh Storage Status")
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: Pure black & white, minimal details, vector-like
- Clean: Simple outlines, minimal shading, easy to color
- Detailed: More intricate lines and fine details
- Bold: Thick outlines, high contrast
- Simple: Basic outlines, perfect for beginners
⚙️ Settings Tips:
- Steps: 15-20 for speed, 25-35 for quality
- Guidance: 7-9 for natural, 10-12 for strict prompt following
- ControlNet: 1.0-1.2 balanced, 1.4+ for strict line following
- Resolution: 768 for speed, 1024+ for detail, 1536 max quality
- Depth Preprocessing: Experimental structural guidance (slower)
""")
# 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],
outputs=[output_image, lineart_preview, depth_preview, status_output, storage_status]
)
clear_cache_btn.click(
manual_clear_cache,
outputs=[status_output, storage_status]
)
refresh_storage_btn.click(
lambda: format_storage_status(),
outputs=storage_status
)
# 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 SDXL - Storage Efficient LineArt")
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")
# 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
)