Skip to content

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

Implement image generation fallback chain: try torch/diffusers first, then sd.cpp 

- Reordered the image generation backend priority to try torch/diffusers first
- If torch/diffusers fails (ImportError or other error), fallback to stable-diffusion-cpp-python
- If both backends fail, return a helpful error message with installation instructions
- Added dynamic loading of sd.cpp model if not pre-loaded
parent 3b527c5a
Show whitespace changes
Inline Side-by-side
Showing with 187 additions and 88 deletions
coderai
View file @ 7f5bf82d
......@@ -3122,8 +3122,115 @@ async def create_image_generation(request: ImageGenerationRequest):
if model_to_use.startswith("image:"):
model_to_use = image_model
# First, try to use stable-diffusion-cpp-python (sd.cpp) if available
# Check all available image models to find one loaded via sd.cpp
# Track errors for proper fallback chain
diffusers_error = None
sd_cpp_error = None
# Parse size (e.g., "1024x1024")
width, height = 1024, 1024
if request.size:
parts = request.size.split("x")
if len(parts) == 2:
try:
width = int(parts[0])
height = int(parts[1])
except ValueError:
pass
# Try diffusers first (torch-based, best quality for NVIDIA)
try:
import torch
from diffusers import StableDiffusionXLPipeline, DiffusionPipeline
# Determine model key
model_key = f"image:{model_to_use}"
pipeline = multi_model_manager.get_model(model_key)
if pipeline is None:
print(f"Loading Stable Diffusion model: {model_to_use}")
# Try to load as Stable Diffusion XL first
try:
pipeline = StableDiffusionXLPipeline.from_pretrained(
model_to_use,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
use_safetensors=True,
)
except Exception:
# Try generic diffusion pipeline
pipeline = DiffusionPipeline.from_pretrained(
model_to_use,
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")
else:
pipeline = pipeline.to("cpu")
# Enable attention slicing for lower memory usage
if torch.cuda.is_available():
pipeline.enable_attention_slicing()
multi_model_manager.add_model(model_key, pipeline)
# Generate images
generator = None
if request.seed is not None:
generator = torch.Generator(device=pipeline.device).manual_seed(request.seed)
# Quality: "standard" or "hd"
quality = request.quality or "standard"
# Generate
result = pipeline(
prompt=request.prompt,
negative_prompt=None,
num_images_per_prompt=request.n,
height=height,
width=width,
generator=generator,
guidance_scale=7.5 if quality == "standard" else 9.0,
num_inference_steps=30 if quality == "standard" else 50,
)
# Extract images
images = []
for img in result.images:
# Convert to base64
import base64
import io
buffered = io.BytesIO()
img.save(buffered, format="PNG")
img_bytes = buffered.getvalue()
img_base64 = base64.b64encode(img_bytes).decode('utf-8')
if request.response_format == "base64":
images.append({"b64_json": img_base64})
else:
# For URL format, we'd need to save somewhere
# For now, return base64
images.append({"b64_json": img_base64})
return {
"created": int(time.time()),
"data": images
}
except ImportError as e:
# diffusers/torch not installed - record error and try sd.cpp
diffusers_error = str(e)
print(f"diffusers not available: {diffusers_error}, trying stable-diffusion-cpp-python...")
except Exception as e:
# Other error with diffusers - record and try sd.cpp
diffusers_error = str(e)
print(f"diffusers error: {diffusers_error}, trying stable-diffusion-cpp-python...")
# Try stable-diffusion-cpp-python (sd.cpp) as fallback
# First, check all available image models to find one loaded via sd.cpp
sd_model = None
for key in multi_model_manager.models:
if key.startswith("image:"):
......@@ -3195,14 +3302,46 @@ async def create_image_generation(request: ImageGenerationRequest):
"created": int(time.time()),
"data": images
}
except ImportError:
pass # stable-diffusion-cpp not available, continue to diffusers
except ImportError as e:
# stable-diffusion-cpp not available
sd_cpp_error = str(e)
print(f"stable-diffusion-cpp-python not available: {sd_cpp_error}")
except Exception as e:
print(f"sd.cpp generation error: {e}")
# Continue to try diffusers
sd_cpp_error = str(e)
else:
# No sd.cpp model pre-loaded, try to load dynamically
print("No pre-loaded sd.cpp model found, trying to load...")
try:
from stable_diffusion_cpp import StableDiffusion
# Parse size (e.g., "1024x1024")
width, height = 1024, 1024
# Check if model_to_use is a URL and get cached path
model_path = None
if model_to_use.startswith('http://') or model_to_use.startswith('https://'):
cached_path = get_cached_model_path(model_to_use)
if cached_path:
model_path = cached_path
print(f"Using cached model: {model_path}")
if model_path is None and os.path.isfile(model_to_use):
model_path = model_to_use
if model_path is None:
print("Warning: Could not resolve sd.cpp model path")
sd_cpp_error = "Could not resolve model path"
else:
# Load sd.cpp model
sd_model = StableDiffusion(
model_path=model_path,
vae_path=None,
n_threads=4,
n_gpu_layers=-1, # All layers to GPU
)
print(f"Using stable-diffusion-cpp-python for image generation")
# Generate images
width, height = 512, 512
if request.size:
parts = request.size.split("x")
if len(parts) == 2:
......@@ -3212,100 +3351,60 @@ async def create_image_generation(request: ImageGenerationRequest):
except ValueError:
pass
# Try to use diffusers if available
try:
import torch
from diffusers import StableDiffusionXLPipeline, DiffusionPipeline
# Determine model key
model_key = f"image:{model_to_use}"
pipeline = multi_model_manager.get_model(model_key)
if pipeline is None:
print(f"Loading Stable Diffusion model: {model_to_use}")
# Try to load as Stable Diffusion XL first
try:
pipeline = StableDiffusionXLPipeline.from_pretrained(
model_to_use,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
use_safetensors=True,
)
except Exception:
# Try generic diffusion pipeline
pipeline = DiffusionPipeline.from_pretrained(
model_to_use,
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")
else:
pipeline = pipeline.to("cpu")
# Enable attention slicing for lower memory usage
if torch.cuda.is_available():
pipeline.enable_attention_slicing()
multi_model_manager.add_model(model_key, pipeline)
# Generate images
generator = None
if request.seed is not None:
generator = torch.Generator(device=pipeline.device).manual_seed(request.seed)
# Quality: "standard" or "hd"
quality = request.quality or "standard"
steps = 4
# Generate
result = pipeline(
result = await asyncio.to_thread(
sd_model.generate_image,
prompt=request.prompt,
negative_prompt=None,
num_images_per_prompt=request.n,
height=height,
negative_prompt='',
width=width,
generator=generator,
guidance_scale=7.5 if quality == "standard" else 9.0,
num_inference_steps=30 if quality == "standard" else 50,
height=height,
cfg_scale=7.0,
sample_steps=steps,
seed=42,
batch_count=request.n if request.n else 1,
)
# Extract images
# Convert results to response format
images = []
for img in result.images:
# Convert to base64
import base64
import io
from PIL import Image
for img in result:
buffered = io.BytesIO()
if isinstance(img, Image.Image):
img.save(buffered, format="PNG")
else:
Image.fromarray(img).save(buffered, format="PNG")
img_bytes = buffered.getvalue()
img_base64 = base64.b64encode(img_bytes).decode('utf-8')
if request.response_format == "base64":
images.append({"b64_json": img_base64})
else:
# For URL format, we'd need to save somewhere
# For now, return base64
images.append({"b64_json": img_base64})
return {
"created": int(time.time()),
"data": images
}
except ImportError as e:
# diffusers not installed
sd_cpp_error = str(e)
print(f"stable-diffusion-cpp-python not available: {sd_cpp_error}")
except Exception as e:
sd_cpp_error = str(e)
print(f"sd.cpp error: {sd_cpp_error}")
# Both backends failed - return error with installation instructions
error_details = []
if diffusers_error:
error_details.append(f"diffusers: {diffusers_error}")
if sd_cpp_error:
error_details.append(f"sd.cpp: {sd_cpp_error}")
raise HTTPException(
status_code=501,
detail=f"Image generation not available. Install diffusers: pip install diffusers torch accelerate safetensors. Error: {str(e)}"
detail=f"Image generation not available. Tried: {', '.join(error_details)}. "
f"Install either: pip install diffusers torch accelerate safetensors (for NVIDIA) "
f"or: pip install stable-diffusion-cpp-python (for Vulkan/AMD)"
)
except Exception as e:
print(f"Image generation error: {e}")
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=f"Image generation error: {str(e)}")
# =============================================================================
# Text-to-Speech Endpoint
# =============================================================================
......
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