Commit 2308d5b0 authored by Your Name's avatar Your Name

Fix image model preloading with --loadall flag

- Fixed bug where image model wasn't actually being loaded when --loadall was specified
- The code only printed messages but never loaded the diffusers pipeline
- Now actually loads the Stable Diffusion pipeline using diffusers library
- Tries StableDiffusionXLPipeline first, falls back to generic DiffusionPipeline
- Moves to GPU if CUDA available, enables attention slicing for memory efficiency
- Also fixes second location where image model is the only configured model

- Debug command line output was already implemented
parent 9193536a
......@@ -4398,6 +4398,44 @@ def main():
# Load image model
if image_models:
print(f"Pre-loading image model: {image_models[0]}")
# Actually load the image model using diffusers
try:
import torch
from diffusers import StableDiffusionXLPipeline, DiffusionPipeline
model_key = f"image:{image_models[0]}"
print(f"Loading diffusers pipeline: {image_models[0]}")
# Try to load as Stable Diffusion XL first
try:
pipeline = StableDiffusionXLPipeline.from_pretrained(
image_models[0],
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
use_safetensors=True,
)
except Exception as e:
print(f"SDXL failed, trying generic pipeline: {e}")
# Try generic diffusion pipeline
pipeline = DiffusionPipeline.from_pretrained(
image_models[0],
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
use_safetensors=True,
)
# Move to GPU if available
if torch.cuda.is_available():
pipeline = pipeline.to("cuda")
pipeline.enable_attention_slicing()
else:
pipeline = pipeline.to("cpu")
multi_model_manager.add_model(model_key, pipeline)
print(f"Image model loaded successfully: {image_models[0]}")
except ImportError as e:
print(f"Warning: diffusers not installed, image model will load on first request: {e}")
except Exception as e:
print(f"Warning: Failed to pre-load image model: {e}")
print(f" Image model will load on first request")
# Load audio model
......@@ -4658,6 +4696,44 @@ def main():
# Pre-load image model if it's the only model configured
if not model_names and not audio_models and not args.tts_model:
print(f"Pre-loading image model...")
# Actually load the image model using diffusers
try:
import torch
from diffusers import StableDiffusionXLPipeline, DiffusionPipeline
model_key = f"image:{image_models[0]}"
print(f"Loading diffusers pipeline: {image_models[0]}")
# Try to load as Stable Diffusion XL first
try:
pipeline = StableDiffusionXLPipeline.from_pretrained(
image_models[0],
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
use_safetensors=True,
)
except Exception as e:
print(f"SDXL failed, trying generic pipeline: {e}")
# Try generic diffusion pipeline
pipeline = DiffusionPipeline.from_pretrained(
image_models[0],
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
use_safetensors=True,
)
# Move to GPU if available
if torch.cuda.is_available():
pipeline = pipeline.to("cuda")
pipeline.enable_attention_slicing()
else:
pipeline = pipeline.to("cpu")
multi_model_manager.add_model(model_key, pipeline)
print(f"Image model loaded successfully: {image_models[0]}")
except ImportError as e:
print(f"Warning: diffusers not installed: {e}")
except Exception as e:
print(f"Warning: Failed to pre-load image model: {e}")
# Register model aliases if specified
if args.model_aliases:
......
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