import torch
from diffusers import StableDiffusionPipeline, DiffusionPipeline

def test_sdg_minimal():
    sdg = DiffusionPipeline.from_pretrained(
        "your-org/safe-diffusion-guidance",
        custom_pipeline="safe_diffusion_guidance",
        torch_dtype=torch.float16
    )
    sdg = sdg.to("cuda" if torch.cuda.is_available() else "cpu")
    out = sdg(
        prompt="test scene",
        base_model_id="runwayml/stable-diffusion-v1-5",
        num_inference_steps=2,  # small for CI
        guidance_scale=5.0,
        safety_scale=2.0,
        mid_fraction=0.5,
        safe_class_index=3
    )
    assert len(out.images) == 1
    print("OK: pipeline end-to-end")
    
if __name__ == "__main__":
    test_sdg_minimal()