test_pipeline.py

"""
Simple script to test if our fixes for diffusers compatibility are working.
This script doesn't use Gradio or the full web interface.
"""

import os
import torch
import numpy as np
from PIL import Image

# Import our custom components
from unet2d_custom import UNet2DConditionModel
from pipeline_stable_diffusion_custom import StableDiffusionPipeline

def main():
    print("Testing SonicDiffusion pipeline components...")
    
    # Check imports
    print("Imports successful!")
    
    # Check if CUDA is available
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    
    # Try to initialize a pipeline (without loading weights, just to test the class structure)
    try:
        # This will just test if the pipeline can be initialized, not if it works correctly
        print("Testing pipeline initialization...")
        pipeline = StableDiffusionPipeline(
            vae=None,
            text_encoder=None,
            tokenizer=None,
            unet=None,
            scheduler=None,
            safety_checker=None,
            feature_extractor=None,
        )
        print("Pipeline initialization successful!")
    except Exception as e:
        print(f"Error initializing pipeline: {e}")
    
    print("Tests completed.")

if __name__ == "__main__":
    main()