Add files using upload-large-folder tool

README.md

@@ -10,6 +10,8 @@ tags:
 pipeline_tag: image-to-image
 ---
 
+> [!WARNING] we do not have a full checkpoint conversion validation, if you encounter pipeline loading failure and unsidered output, please contact me via bili_sakura@zju.edu.cn
+
 # BiliSakura/HSIGene
 
 **Hyperspectral image generation** — HSIGene converted to diffusers format. Supports task-specific conditioning with local controls (HED, MLSD, sketch, segmentation), global controls (content or text), or metadata embeddings. Outputs 48-band hyperspectral images (256×256 pixels).

pipeline_hsigene.py

@@ -139,7 +139,10 @@ class _CRSModelWrapper(torch.nn.Module):
         local_control_scales=None,
     ):
         super().__init__()
-        self.model = type("Model", (), {"diffusion_model": unet})()
+        # Keep diffusion_model as a properly registered submodule so
+        # wrapper/device transfers (e.g., `.to("cuda")`) move UNet weights.
+        self.model = torch.nn.Module()
+        self.model.add_module("diffusion_model", unet)
         self.first_stage_model = vae
         self.cond_stage_model = text_encoder
         self.local_adapter = local_adapter
@@ -408,7 +411,18 @@ class HSIGenePipeline(DiffusionPipeline):
         return_dict: bool = True,
         save_memory: bool = False,
     ):
-        device = next(self.crs_model.parameters()).device
+        target_device = next(self.crs_model.parameters()).device
+        if hasattr(self, "unet") and isinstance(self.unet, torch.nn.Module):
+            target_device = next(self.unet.parameters()).device
+        if latents is not None:
+            target_device = latents.device
+        elif generator is not None and hasattr(generator, "device"):
+            target_device = torch.device(generator.device)
+
+        # Keep wrapper submodules on the same device used for sampling.
+        if next(self.crs_model.parameters()).device != target_device:
+            self.crs_model = self.crs_model.to(target_device)
+        device = target_device
         if text_strength is None:
             text_strength = global_strength
 
@@ -468,6 +482,15 @@ class HSIGenePipeline(DiffusionPipeline):
 
         latent_shape = (num_samples, 4, height // 4, width // 4)
         if latents is None:
+            if generator is not None and hasattr(generator, "device"):
+                gen_device = torch.device(generator.device)
+                if gen_device.type != device.type:
+                    # Recreate generator on target device while preserving seed
+                    # so CPU/CUDA mismatch does not crash torch.randn.
+                    if hasattr(generator, "initial_seed"):
+                        generator = torch.Generator(device=device).manual_seed(generator.initial_seed())
+                    else:
+                        generator = torch.Generator(device=device)
             latents = torch.randn(
                 latent_shape, device=device, generator=generator, dtype=torch.float32,
             )