Fix SD15 dtype handling off CUDA

ip_adapter/ip_adapter.py

@@ -71,11 +71,13 @@ class IPAdapter:
         self.num_tokens = num_tokens
 
         self.pipe = sd_pipe.to(self.device)
+        self.torch_dtype = next(self.pipe.unet.parameters()).dtype
         self.set_ip_adapter()
 
         # load image encoder
         self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(self.image_encoder_path).to(
-            self.device, dtype=torch.float16
+            self.device,
+            dtype=self.torch_dtype,
         )
         self.clip_image_processor = CLIPImageProcessor()
         # image proj model
@@ -88,7 +90,7 @@ class IPAdapter:
             cross_attention_dim=self.pipe.unet.config.cross_attention_dim,
             clip_embeddings_dim=self.image_encoder.config.projection_dim,
             clip_extra_context_tokens=self.num_tokens,
-        ).to(self.device, dtype=torch.float16)
+        ).to(self.device, dtype=self.torch_dtype)
         return image_proj_model
 
     def set_ip_adapter(self):
@@ -112,7 +114,7 @@ class IPAdapter:
                     cross_attention_dim=cross_attention_dim,
                     scale=1.0,
                     num_tokens=self.num_tokens,
-                ).to(self.device, dtype=torch.float16)
+                ).to(self.device, dtype=self.torch_dtype)
         unet.set_attn_processor(attn_procs)
         if hasattr(self.pipe, "controlnet"):
             if isinstance(self.pipe.controlnet, MultiControlNetModel):
@@ -142,9 +144,11 @@ class IPAdapter:
             if isinstance(pil_image, Image.Image):
                 pil_image = [pil_image]
             clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
-            clip_image_embeds = self.image_encoder(clip_image.to(self.device, dtype=torch.float16)).image_embeds
+            clip_image_embeds = self.image_encoder(
+                clip_image.to(self.device, dtype=self.torch_dtype)
+            ).image_embeds
         else:
-            clip_image_embeds = clip_image_embeds.to(self.device, dtype=torch.float16)
+            clip_image_embeds = clip_image_embeds.to(self.device, dtype=self.torch_dtype)
         image_prompt_embeds = self.image_proj_model(clip_image_embeds)
         uncond_image_prompt_embeds = self.image_proj_model(torch.zeros_like(clip_image_embeds))
         return image_prompt_embeds, uncond_image_prompt_embeds
@@ -296,16 +300,29 @@ class IPAdapterPlus(IPAdapter):
             embedding_dim=self.image_encoder.config.hidden_size,
             output_dim=self.pipe.unet.config.cross_attention_dim,
             ff_mult=4,
-        ).to(self.device, dtype=torch.float16)
+        ).to(self.device, dtype=self.torch_dtype)
         return image_proj_model
 
     @torch.inference_mode()
     def get_image_embeds(self, pil_image=None, clip_image_embeds=None):
-        if isinstance(pil_image, Image.Image):
-            pil_image = [pil_image]
-        clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
-        clip_image = clip_image.to(self.device, dtype=torch.float16)
-        clip_image_embeds = self.image_encoder(clip_image, output_hidden_states=True).hidden_states[-2]
+        if pil_image is not None:
+            if isinstance(pil_image, Image.Image):
+                pil_image = [pil_image]
+            clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
+            clip_image = clip_image.to(self.device, dtype=self.torch_dtype)
+            clip_image_embeds = self.image_encoder(
+                clip_image, output_hidden_states=True
+            ).hidden_states[-2]
+        else:
+            clip_image_embeds = clip_image_embeds.to(self.device, dtype=self.torch_dtype)
+            clip_image = torch.zeros(
+                clip_image_embeds.shape[0],
+                3,
+                224,
+                224,
+                device=self.device,
+                dtype=self.torch_dtype,
+            )
         image_prompt_embeds = self.image_proj_model(clip_image_embeds)
         uncond_clip_image_embeds = self.image_encoder(
             torch.zeros_like(clip_image), output_hidden_states=True
@@ -321,7 +338,7 @@ class IPAdapterFull(IPAdapterPlus):
         image_proj_model = MLPProjModel(
             cross_attention_dim=self.pipe.unet.config.cross_attention_dim,
             clip_embeddings_dim=self.image_encoder.config.hidden_size,
-        ).to(self.device, dtype=torch.float16)
+        ).to(self.device, dtype=self.torch_dtype)
         return image_proj_model
 
 

ipadapter_model.py

@@ -8,6 +8,8 @@ This module provides utilities for working with IP-Adapter models, including:
 - Utility functions for image processing
 """
 
+import logging
+import os
 from typing import List, Optional, Union, Tuple
 
 import numpy as np
@@ -21,6 +23,91 @@ torch.backends.cuda.enable_cudnn_sdp(False)
 from ip_adapter import IPAdapterPlus, IPAdapterPlusXL
 
 
+# ===== SD1.5 Base Model Resolution =====
+
+DEFAULT_SD15_BASE_MODEL = "SG161222/Realistic_Vision_V4.0_noVAE"
+DEFAULT_SD15_FALLBACK_MODELS = (
+    "runwayml/stable-diffusion-v1-5",
+    "stable-diffusion-v1-5/stable-diffusion-v1-5",
+)
+SD15_BASE_MODEL_ENV = "VIBESPACE_SD15_BASE_MODEL"
+SD15_FALLBACK_MODELS_ENV = "VIBESPACE_SD15_FALLBACK_MODELS"
+
+
+def _get_sd15_torch_dtype(device: str) -> torch.dtype:
+    """Keep SD1.5 in fp16 on CUDA and use fp32 elsewhere for numerical stability."""
+    return torch.float16 if str(device).lower().startswith("cuda") else torch.float32
+
+
+def _get_ip_model_dtype(ip_model) -> torch.dtype:
+    return getattr(ip_model, "torch_dtype", torch.float16)
+
+
+def _get_sd15_base_model_candidates() -> List[str]:
+    """Build an ordered, de-duplicated list of SD1.5 base model candidates."""
+    configured_primary = os.getenv(SD15_BASE_MODEL_ENV, "").strip()
+    configured_fallbacks = os.getenv(SD15_FALLBACK_MODELS_ENV, "").strip()
+
+    candidates: List[str] = []
+    if configured_primary:
+        candidates.append(configured_primary)
+
+    candidates.append(DEFAULT_SD15_BASE_MODEL)
+
+    if configured_fallbacks:
+        candidates.extend(
+            model_id.strip()
+            for model_id in configured_fallbacks.split(",")
+            if model_id.strip()
+        )
+    else:
+        candidates.extend(DEFAULT_SD15_FALLBACK_MODELS)
+
+    deduplicated_candidates: List[str] = []
+    seen = set()
+    for model_id in candidates:
+        if model_id in seen:
+            continue
+        deduplicated_candidates.append(model_id)
+        seen.add(model_id)
+
+    return deduplicated_candidates
+
+
+def _load_sd15_base_pipeline(
+    noise_scheduler: DDIMScheduler,
+    vae: AutoencoderKL,
+    torch_dtype: torch.dtype,
+) -> StableDiffusionPipeline:
+    """Load the first available SD1.5-compatible base pipeline."""
+    candidates = _get_sd15_base_model_candidates()
+    last_error: Optional[Exception] = None
+
+    for index, base_model_path in enumerate(candidates):
+        try:
+            return StableDiffusionPipeline.from_pretrained(
+                base_model_path,
+                torch_dtype=torch_dtype,
+                scheduler=noise_scheduler,
+                vae=vae,
+                feature_extractor=None,
+                safety_checker=None,
+            )
+        except Exception as exc:  # noqa: BLE001
+            last_error = exc
+            if index < len(candidates) - 1:
+                logging.warning(
+                    "Failed to load SD1.5 base model '%s'; trying fallback. Error: %s",
+                    base_model_path,
+                    exc,
+                )
+
+    candidate_list = ", ".join(candidates)
+    raise RuntimeError(
+        f"Failed to load any SD1.5 base model. Tried: {candidate_list}"
+    ) from last_error
+
+
 # ===== Image Utility Functions =====
 
 def create_image_grid(images: List[Image.Image], rows: int, cols: int) -> Image.Image:
@@ -57,7 +144,10 @@ def extract_clip_embeddings_from_pil(pil_image: Union[Image.Image, List[Image.Im
     ).pixel_values
     
     # Move to model device with appropriate dtype
-    processed_images = processed_images.to(ip_model.device, dtype=torch.float16)
+    processed_images = processed_images.to(
+        ip_model.device,
+        dtype=_get_ip_model_dtype(ip_model),
+    )
     
     # Extract embeddings from penultimate layer (better for downstream tasks)
     clip_embeddings = ip_model.image_encoder(
@@ -93,7 +183,10 @@ def extract_clip_embeddings_from_tensor(tensor_image: torch.Tensor,
         torch.Tensor: CLIP embeddings of shape (batch_size, seq_len, embed_dim)
     """
     # Move tensor to model device with appropriate dtype
-    tensor_image = tensor_image.to(ip_model.device, dtype=torch.float16)
+    tensor_image = tensor_image.to(
+        ip_model.device,
+        dtype=_get_ip_model_dtype(ip_model),
+    )
     
     # Resize to CLIP input resolution if requested
     if resize:
@@ -133,6 +226,7 @@ def _enhanced_get_image_embeds(self, pil_image=None, clip_image_embeds=None):
         Tuple of (conditional_embeds, unconditional_embeds)
     """
     # Process PIL images if provided
+    model_dtype = getattr(self, "torch_dtype", torch.float16)
     if pil_image is not None:
         if isinstance(pil_image, Image.Image):
             pil_image = [pil_image]
@@ -141,17 +235,26 @@ def _enhanced_get_image_embeds(self, pil_image=None, clip_image_embeds=None):
         processed_images = self.clip_image_processor(
             images=pil_image, return_tensors="pt"
         ).pixel_values
-        processed_images = processed_images.to(self.device, dtype=torch.float16)
+        processed_images = processed_images.to(self.device, dtype=model_dtype)
         
         clip_image_embeds = self.image_encoder(
             processed_images, output_hidden_states=True
         ).hidden_states[-2]
+    else:
+        clip_image_embeds = clip_image_embeds.to(self.device, dtype=model_dtype)
     
     # Project CLIP embeddings to IP-Adapter space
     conditional_embeds = self.image_proj_model(clip_image_embeds)
     
     # Generate unconditional embeddings (for classifier-free guidance)
-    zero_tensor = torch.zeros(1, 3, 224, 224).to(self.device, dtype=torch.float16)
+    zero_tensor = torch.zeros(
+        clip_image_embeds.shape[0],
+        3,
+        224,
+        224,
+        device=self.device,
+        dtype=model_dtype,
+    )
     uncond_clip_embeds = self.image_encoder(
         zero_tensor, output_hidden_states=True
     ).hidden_states[-2]
@@ -164,9 +267,8 @@ def _enhanced_get_image_embeds(self, pil_image=None, clip_image_embeds=None):
 
 @torch.inference_mode()
 def load_stable_diffusion_pipeline(device: str = "cuda") -> StableDiffusionPipeline:
-    # Model paths
-    base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
     vae_model_path = "stabilityai/sd-vae-ft-mse"
+    torch_dtype = _get_sd15_torch_dtype(device)
 
     # Configure DDIM scheduler for high-quality sampling
     noise_scheduler = DDIMScheduler(
@@ -180,28 +282,28 @@ def load_stable_diffusion_pipeline(device: str = "cuda") -> StableDiffusionPipel
     )
     
     # Load VAE separately for better quality
-    vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
+    vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch_dtype)
     
-    # Create Stable Diffusion pipeline
-    pipeline = StableDiffusionPipeline.from_pretrained(
-        base_model_path,
-        torch_dtype=torch.float16,
-        scheduler=noise_scheduler,
+    # Create Stable Diffusion pipeline with fallback SD1.5 bases
+    pipeline = _load_sd15_base_pipeline(
+        noise_scheduler=noise_scheduler,
         vae=vae,
-        feature_extractor=None,  # Disable safety checker for faster inference
-        safety_checker=None,
+        torch_dtype=torch_dtype,
     )
     
     return pipeline
 
 
 @torch.inference_mode()
-def load_ip_adapter_model(device: str = "cuda", sd_only: bool = False) -> IPAdapterPlus:
+def load_ip_adapter_model(
+    device: str = "cuda",
+    sd_only: bool = False,
+) -> IPAdapterPlus | StableDiffusionPipeline:
     # Model and checkpoint paths
-    base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
     vae_model_path = "stabilityai/sd-vae-ft-mse"
     image_encoder_path = "./downloads/models/image_encoder"
     ip_checkpoint_path = "./downloads/models/ip-adapter-plus_sd15.bin"
+    torch_dtype = _get_sd15_torch_dtype(device)
 
     # Configure DDIM scheduler
     noise_scheduler = DDIMScheduler(
@@ -215,16 +317,13 @@ def load_ip_adapter_model(device: str = "cuda", sd_only: bool = False) -> IPAdap
     )
     
     # Load high-quality VAE
-    vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
+    vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch_dtype)
     
-    # Create base Stable Diffusion pipeline
-    pipeline = StableDiffusionPipeline.from_pretrained(
-        base_model_path,
-        torch_dtype=torch.float16,
-        scheduler=noise_scheduler,
+    # Create base Stable Diffusion pipeline with fallback SD1.5 bases
+    pipeline = _load_sd15_base_pipeline(
+        noise_scheduler=noise_scheduler,
         vae=vae,
-        feature_extractor=None,
-        safety_checker=None,
+        torch_dtype=torch_dtype,
     )
     
     if sd_only:
@@ -287,7 +386,10 @@ def generate_images_from_clip_embeddings(ip_model : IPAdapterPlus,
         raise ValueError(f"Expected 3D embeddings (batch, seq, dim), got {clip_embeddings.shape}")
     
     # Move to appropriate device and dtype
-    clip_embeddings = clip_embeddings.half().to(ip_model.device)
+    clip_embeddings = clip_embeddings.to(
+        ip_model.device,
+        dtype=_get_ip_model_dtype(ip_model),
+    )
     
     # Generate images using IP-Adapter
     negative_prompt = "nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]"
@@ -311,4 +413,4 @@ extract_clip_embedding_pil = extract_clip_embeddings_from_pil
 extract_clip_embedding_pil_batch = extract_clip_embeddings_from_pil_batch
 extract_clip_embedding_tensor = extract_clip_embeddings_from_tensor
 load_sdxl = load_stable_diffusion_pipeline
-generate = generate_images_from_clip_embeddings
+generate = generate_images_from_clip_embeddings