DobidopandClaude
c2b3bb5a50
Combines three powerful stereoscopic tools into one comprehensive package: 1. Core Stereo Generation - Depth-based stereo conversion with GPU acceleration - Advanced fill techniques (Polylines, GPU Warp, Hybrid Edge, etc.) - Edge-aware depth blurring for artifact reduction - Side-by-Side, Top-Bottom, and Red-Cyan Anaglyph output formats 2. Native VR Viewing (Optional - requires PyOpenXR) - Auto-launch directly into VR headset without browser - Support for all VR headsets (Quest, Vive, Index, WMR) - Multiple projection types (Flat, Curved, 180° Dome, 360° Sphere) - Image and video playback support 3. StereoDiffusion AI (Optional - requires diffusers) - AI-powered stereo generation using diffusion models - DDIM inversion with null-text optimization - Bilateral Neighbor Attention for stereo consistency - Works with ComfyUI MODEL/CLIP/VAE or HuggingFace models Changes: - Added native_viewer/ module with PyOpenXR implementation - Added native_nodes.py with 3 VR viewer nodes - Added stereodiffusion_nodes.py with AI stereo generation - Added supporting modules: model_wrappers, model_loader, inversion, stereo_utils, diffusion_utils - Updated __init__.py with graceful optional dependency handling - Consolidated all requirements into single requirements.txt - Comprehensive README with setup guides and troubleshooting - Added NOTICE and licenses/ for proper attribution The package now offers 6+ nodes total with modular installation options. Users can install just the base, or add VR viewing and/or AI generation. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
"""
|
|
Model Loading Utilities
|
|
|
|
Functions for loading Stable Diffusion models from HuggingFace or local paths.
|
|
"""
|
|
|
|
import os
|
|
import torch
|
|
from diffusers import StableDiffusionPipeline, DDIMScheduler
|
|
|
|
|
|
# Global model cache
|
|
_model_cache = {}
|
|
|
|
|
|
def clear_model_cache():
|
|
"""Clear the model cache to force reloading."""
|
|
global _model_cache
|
|
_model_cache = {}
|
|
|
|
|
|
def load_sd_model(model_id_or_path: str = "runwayml/stable-diffusion-v1-5", device: str = "cuda"):
|
|
"""
|
|
Load or retrieve cached Stable Diffusion model.
|
|
|
|
Args:
|
|
model_id_or_path: Either a HuggingFace model ID (e.g., "runwayml/stable-diffusion-v1-5")
|
|
or a local path to a diffusers-format model directory.
|
|
device: Device to load the model on ("cuda" or "cpu")
|
|
|
|
Returns:
|
|
StableDiffusionPipeline instance
|
|
"""
|
|
global _model_cache
|
|
|
|
if model_id_or_path in _model_cache:
|
|
return _model_cache[model_id_or_path]
|
|
|
|
scheduler = DDIMScheduler(
|
|
beta_start=0.00085,
|
|
beta_end=0.012,
|
|
beta_schedule="scaled_linear",
|
|
clip_sample=False,
|
|
set_alpha_to_one=False
|
|
)
|
|
|
|
# Check if it's a local path
|
|
is_local = os.path.isdir(model_id_or_path)
|
|
|
|
# Use float32 for null-text optimization (requires gradients)
|
|
# float16 can cause NaN due to limited precision in gradient computations
|
|
try:
|
|
model = StableDiffusionPipeline.from_pretrained(
|
|
model_id_or_path,
|
|
scheduler=scheduler,
|
|
torch_dtype=torch.float32, # Use float32 for gradient stability
|
|
local_files_only=is_local,
|
|
safety_checker=None, # Disable safety checker for faster loading
|
|
requires_safety_checker=False,
|
|
).to(device)
|
|
except Exception as e:
|
|
print(f"Failed to load model: {e}")
|
|
print("Attempting to download from HuggingFace...")
|
|
model = StableDiffusionPipeline.from_pretrained(
|
|
model_id_or_path,
|
|
scheduler=scheduler,
|
|
torch_dtype=torch.float32,
|
|
safety_checker=None,
|
|
requires_safety_checker=False,
|
|
).to(device)
|
|
|
|
try:
|
|
model.disable_xformers_memory_efficient_attention()
|
|
except AttributeError:
|
|
pass
|
|
|
|
_model_cache[model_id_or_path] = model
|
|
return model
|