Update all files for DiffusionSat-SR-Texas-256

pipeline_diffusionsat_controlnet.py

@@ -8,6 +8,13 @@ from __future__ import annotations
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
+# Import types directly (not just for TYPE_CHECKING) so diffusers can introspect them
+from diffusers.models import AutoencoderKL
+from diffusers.models.controlnets import ControlNetModel
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+
 import einops
 import numpy as np
 import PIL.Image
@@ -22,11 +29,11 @@ from diffusers.schedulers import KarrasDiffusionSchedulers
 from diffusers.utils import (
     PIL_INTERPOLATION,
     logging,
-    randn_tensor,
     replace_example_docstring,
     is_accelerate_available,
     is_accelerate_version,
 )
+from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput
 from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
@@ -83,11 +90,11 @@ class DiffusionSatControlNetPipeline(DiffusionPipeline, TextualInversionLoaderMi
         vae: AutoencoderKL,
         text_encoder: CLIPTextModel,
         tokenizer: CLIPTokenizer,
-        unet: Any,
-        controlnet: Any,
+        unet: UNet2DConditionModel,
+        controlnet: ControlNetModel,
         scheduler: KarrasDiffusionSchedulers,
-        safety_checker: StableDiffusionSafetyChecker,
-        feature_extractor: CLIPImageProcessor,
+        safety_checker: Optional[StableDiffusionSafetyChecker] = None,
+        feature_extractor: Optional[CLIPImageProcessor] = None,
         requires_safety_checker: bool = True,
     ):
         super().__init__()
@@ -229,7 +236,12 @@ class DiffusionSatControlNetPipeline(DiffusionPipeline, TextualInversionLoaderMi
         cond_metadata: Optional[List[float]] = None,
         is_temporal: bool = False,
         conditioning_downsample: bool = True,
-    ):
+    ) -> Union[StableDiffusionPipelineOutput, Tuple]:
+        """
+        Function invoked when calling the pipeline for generation.
+        
+        Examples:
+        """
         # 0. Default height and width to unet
         height, width = self._default_height_width(height, width, image)
         cond_height, cond_width = height, width