1+ using System . Security . Cryptography ;
2+ using FluentAssertions ;
3+ using FluentAssertions . Execution ;
4+ using Microsoft . Extensions . DependencyInjection ;
5+ using Microsoft . Extensions . Logging ;
6+ using OnnxStack . Core ;
7+ using OnnxStack . StableDiffusion . Common ;
8+ using OnnxStack . StableDiffusion . Config ;
9+ using OnnxStack . StableDiffusion . Enums ;
10+ using SixLabors . ImageSharp ;
11+ using Xunit . Abstractions ;
12+
13+ namespace OnnxStack . IntegrationTests ;
14+
15+ /// <summary>
16+ /// These tests just run on CPU execution provider for now, but could switch it to CUDA and run on GPU
17+ /// if the necessary work is done to setup the docker container to allow GPU passthrough to the container.
18+ /// See https://blog.roboflow.com/use-the-gpu-in-docker/ for an example of how to do this.
19+ ///
20+ /// Can then also setup a self-hosted runner in Github Actions to run the tests on your own GPU as part of the CI/CD pipeline.
21+ /// Maybe something like https://www.youtube.com/watch?v=rVq-SCNyxVc
22+ /// </summary>
23+ public class StableDiffusionTests
24+ {
25+ private readonly IStableDiffusionService _stableDiffusion ;
26+ private readonly ILogger < StableDiffusionTests > _logger ;
27+ private const string StableDiffusionModel = "StableDiffusion 1.5" ;
28+ private const string LatentConsistencyModel = "LCM-Dreamshaper-V7" ;
29+
30+ public StableDiffusionTests ( ITestOutputHelper testOutputHelper )
31+ {
32+ var services = new ServiceCollection ( ) ;
33+ services . AddLogging ( builder => builder . AddConsole ( ) ) ; //necessary for showing logs when running in docker
34+ services . AddLogging ( builder => builder . AddXunit ( testOutputHelper ) ) ; //necessary for showing logs when running in IDE
35+ services . AddOnnxStackStableDiffusion ( ) ;
36+ var provider = services . BuildServiceProvider ( ) ;
37+ _stableDiffusion = provider . GetRequiredService < IStableDiffusionService > ( ) ;
38+ _logger = provider . GetRequiredService < ILogger < StableDiffusionTests > > ( ) ;
39+ }
40+
41+ [ Theory ]
42+ [ InlineData ( StableDiffusionModel ) ]
43+ [ InlineData ( LatentConsistencyModel ) ]
44+ public async Task GivenAStableDiffusionModel_WhenLoadModel_ThenModelIsLoaded ( string modelName )
45+ {
46+ //arrange
47+ var model = _stableDiffusion . Models . Single ( m => m . Name == modelName ) ;
48+
49+ //act
50+ _logger . LogInformation ( "Attempting to load model {0}" , model . Name ) ;
51+ var isModelLoaded = await _stableDiffusion . LoadModel ( model ) ;
52+
53+ //assert
54+ isModelLoaded . Should ( ) . BeTrue ( ) ;
55+ }
56+
57+ [ Theory ]
58+ [ InlineData ( StableDiffusionModel , SchedulerType . EulerAncestral , 10 , 7.0f , "E518D0E4F67CBD5E93513574D30F3FD7" ) ]
59+ [ InlineData ( LatentConsistencyModel , SchedulerType . LCM , 4 , 1.0f , "3554E5E1B714D936805F4C9D890B0711" ) ]
60+ public async Task GivenTextToImage_WhenInference_ThenImageGenerated ( string modelName , SchedulerType schedulerType ,
61+ int inferenceSteps , float guidanceScale , string generatedImageMd5Hash )
62+
63+ {
64+ //arrange
65+ var model = _stableDiffusion . Models . Single ( m => m . Name == modelName ) ;
66+ _logger . LogInformation ( "Attempting to load model: {0}" , model . Name ) ;
67+ await _stableDiffusion . LoadModel ( model ) ;
68+
69+ var prompt = new PromptOptions
70+ {
71+ Prompt = "an astronaut riding a horse in space" ,
72+ NegativePrompt = "blurry,ugly,cartoon" ,
73+ BatchCount = 1 ,
74+ DiffuserType = DiffuserType . TextToImage
75+ } ;
76+
77+ var scheduler = new SchedulerOptions
78+ {
79+ Width = 512 ,
80+ Height = 512 ,
81+ SchedulerType = schedulerType ,
82+ InferenceSteps = inferenceSteps ,
83+ GuidanceScale = guidanceScale ,
84+ Seed = 1
85+ } ;
86+
87+ var steps = 0 ;
88+
89+ //act
90+ var image = await _stableDiffusion . GenerateAsImageAsync ( model , prompt , scheduler , ( currentStep , totalSteps ) =>
91+ {
92+ _logger . LogInformation ( $ "Step { currentStep } /{ totalSteps } ") ;
93+ steps ++ ;
94+ } ) ;
95+
96+ var imagesDirectory = Path . Combine ( Directory . GetCurrentDirectory ( ) , "images" ) ;
97+ if ( ! Directory . Exists ( imagesDirectory ) )
98+ {
99+ _logger . LogInformation ( $ "Creating directory { imagesDirectory } ") ;
100+ Directory . CreateDirectory ( imagesDirectory ) ;
101+ }
102+ else
103+ {
104+ _logger . LogInformation ( $ "Directory { imagesDirectory } already exists") ;
105+ }
106+
107+ var fileName =
108+ $ "{ imagesDirectory } /{ nameof ( GivenTextToImage_WhenInference_ThenImageGenerated ) } -{ DateTime . Now : yyyyMMddHHmmss} .png";
109+ _logger . LogInformation ( $ "Saving generated image to { fileName } ") ;
110+ await image . SaveAsPngAsync ( fileName ) ;
111+
112+ //assert
113+ using ( new AssertionScope ( ) )
114+ {
115+ steps . Should ( ) . Be ( inferenceSteps ) ;
116+ image . Should ( ) . NotBeNull ( ) ;
117+ image . Size . IsEmpty . Should ( ) . BeFalse ( ) ;
118+ image . Width . Should ( ) . Be ( 512 ) ;
119+ image . Height . Should ( ) . Be ( 512 ) ;
120+
121+ File . Exists ( fileName ) . Should ( ) . BeTrue ( ) ;
122+ var md5 = MD5 . Create ( ) ;
123+ var hash = md5 . ComputeHash ( File . ReadAllBytes ( fileName ) ) ;
124+ var hashString = string . Join ( "" , hash . Select ( b => b . ToString ( "X2" ) ) ) ;
125+ _logger . LogInformation ( $ "MD5 Hash of generated image: { hashString } ") ;
126+
127+ hashString . Should ( ) . Be ( generatedImageMd5Hash ) ;
128+ }
129+ }
130+ }
0 commit comments