11using Microsoft . AspNetCore . SignalR ;
22using OnnxStack . StableDiffusion . Common ;
33using OnnxStack . StableDiffusion . Config ;
4- using OnnxStack . StableDiffusion . Models ;
54using OnnxStack . WebUI . Models ;
65using Services ;
76using SixLabors . ImageSharp ;
@@ -96,6 +95,62 @@ public async IAsyncEnumerable<StableDiffusionResult> OnExecuteImageToImage(Promp
9695 }
9796
9897
98+ /// <summary>
99+ /// Execute Image-Inpaint Stable Diffusion
100+ /// </summary>
101+ /// <param name="options">The options.</param>
102+ /// <param name="cancellationToken">The cancellation token.</param>
103+ /// <returns></returns>
104+ [ HubMethodName ( "ExecuteImageInpaint" ) ]
105+ public async IAsyncEnumerable < StableDiffusionResult > OnExecuteImageInpaint ( PromptOptions promptOptions , SchedulerOptions schedulerOptions , [ EnumeratorCancellation ] CancellationToken cancellationToken )
106+ {
107+ _logger . Log ( LogLevel . Information , "[ExecuteImageInpaint] - New request received, Connection: {0}" , Context . ConnectionId ) ;
108+ var cancellationTokenSource = CancellationTokenSource . CreateLinkedTokenSource ( Context . ConnectionAborted , cancellationToken ) ;
109+
110+ // TODO: Add support for multiple results
111+ var result = await GenerateImageInpaintResult ( promptOptions , schedulerOptions , cancellationTokenSource . Token ) ;
112+ if ( ! result . IsError )
113+ yield return result ;
114+
115+ await Clients . Caller . OnError ( result . Error ) ;
116+ yield break ;
117+ }
118+
119+
120+ /// <summary>
121+ /// Generates the text to image result.
122+ /// </summary>
123+ /// <param name="options">The options.</param>
124+ /// <param name="cancellationToken">The cancellation token.</param>
125+ /// <returns></returns>
126+ private async Task < StableDiffusionResult > GenerateTextToImageResult ( PromptOptions promptOptions , SchedulerOptions schedulerOptions , CancellationToken cancellationToken )
127+ {
128+ var timestamp = Stopwatch . GetTimestamp ( ) ;
129+ schedulerOptions . Seed = GenerateSeed ( schedulerOptions . Seed ) ;
130+
131+ //1. Create filenames
132+ var random = await _fileService . CreateRandomName ( ) ;
133+ var outputImage = $ "{ random } -output.png";
134+ var outputBlueprint = $ "{ random } -output.json";
135+ var outputImageUrl = await _fileService . CreateOutputUrl ( outputImage ) ;
136+ var outputImageFile = await _fileService . UrlToPhysicalPath ( outputImageUrl ) ;
137+ var outputImageLink = await _fileService . CreateOutputUrl ( outputImage , false ) ;
138+
139+ //2. Generate blueprint
140+ var blueprint = new ImageBlueprint ( promptOptions , schedulerOptions , outputImageLink ) ;
141+ var bluprintFile = await _fileService . SaveBlueprintFile ( blueprint , outputBlueprint ) ;
142+ if ( bluprintFile is null )
143+ return new StableDiffusionResult ( "Failed to save blueprint" ) ;
144+
145+ //3. Run stable diffusion
146+ if ( ! await RunStableDiffusion ( promptOptions , schedulerOptions , outputImageFile , cancellationToken ) )
147+ return new StableDiffusionResult ( "Failed to run stable diffusion" ) ;
148+
149+ //4. Return result
150+ return new StableDiffusionResult ( outputImage , outputImageUrl , blueprint , bluprintFile . Filename , bluprintFile . FileUrl , GetElapsed ( timestamp ) ) ;
151+ }
152+
153+
99154 /// <summary>
100155 /// Generates the image to image result.
101156 /// </summary>
@@ -110,72 +165,79 @@ private async Task<StableDiffusionResult> GenerateImageToImageResult(PromptOptio
110165
111166 //1. Create filenames
112167 var random = await _fileService . CreateRandomName ( ) ;
113- var output = $ "Output-{ random } ";
114- var outputImage = $ "{ output } .png";
115- var outputBlueprint = $ "{ output } .json";
116- var inputImage = $ "Input-{ random } .png";
117- var uploadImage = Path . GetFileName ( promptOptions . InputImage . ImagePath ) ;
168+ var inputImage = $ "{ random } -input.png";
169+ var outputImage = $ "{ random } -output.png";
170+ var outputBlueprint = $ "{ random } -output.json";
118171 var outputImageUrl = await _fileService . CreateOutputUrl ( outputImage ) ;
119172 var outputImageFile = await _fileService . UrlToPhysicalPath ( outputImageUrl ) ;
120- var inputOriginaUrl = await _fileService . CreateOutputUrl ( uploadImage , false ) ;
173+ var outputImageLink = await _fileService . CreateOutputUrl ( outputImage , false ) ;
174+ var inputImageLink = await _fileService . CreateOutputUrl ( inputImage , false ) ;
121175
122- //2. Copy input image to new file
123- var inputImageFile = await _fileService . CopyInputImageFile ( promptOptions . InputImage . ImagePath , inputImage ) ;
176+ // 2. Save Input Image to disk
177+ var inputImageFile = await _fileService . SaveImageFile ( promptOptions . InputImage . ImageBase64 , inputImage ) ;
124178 if ( inputImageFile is null )
125- return new StableDiffusionResult ( "Failed to copy input image" ) ;
179+ return new StableDiffusionResult ( "Failed to save input image" ) ;
126180
127- //3. Generate blueprint
128- var inputImageLink = await _fileService . CreateOutputUrl ( inputImage , false ) ;
129- var outputImageLink = await _fileService . CreateOutputUrl ( outputImage , false ) ;
130- promptOptions . InputImage = new InputImage ( inputOriginaUrl ) ;
181+ //3. Generate and save blueprint file
131182 var blueprint = new ImageBlueprint ( promptOptions , schedulerOptions , outputImageLink , inputImageLink ) ;
132183 var bluprintFile = await _fileService . SaveBlueprintFile ( blueprint , outputBlueprint ) ;
133184 if ( bluprintFile is null )
134185 return new StableDiffusionResult ( "Failed to save blueprint" ) ;
135186
136- //4. Set full path of input image
137- promptOptions . InputImage = new InputImage ( inputImageFile . FilePath ) ;
138-
139- //5. Run stable diffusion
187+ //4. Run stable diffusion
140188 if ( ! await RunStableDiffusion ( promptOptions , schedulerOptions , outputImageFile , cancellationToken ) )
141189 return new StableDiffusionResult ( "Failed to run stable diffusion" ) ;
142190
143- //6 . Return result
191+ //5 . Return result
144192 return new StableDiffusionResult ( outputImage , outputImageUrl , blueprint , bluprintFile . Filename , bluprintFile . FileUrl , GetElapsed ( timestamp ) ) ;
145193 }
146194
147195
148196 /// <summary>
149- /// Generates the text to image result.
197+ /// Generates the image inpaint result.
150198 /// </summary>
151- /// <param name="options">The options.</param>
199+ /// <param name="promptOptions">The prompt options.</param>
200+ /// <param name="schedulerOptions">The scheduler options.</param>
152201 /// <param name="cancellationToken">The cancellation token.</param>
153202 /// <returns></returns>
154- private async Task < StableDiffusionResult > GenerateTextToImageResult ( PromptOptions promptOptions , SchedulerOptions schedulerOptions , CancellationToken cancellationToken )
203+ private async Task < StableDiffusionResult > GenerateImageInpaintResult ( PromptOptions promptOptions , SchedulerOptions schedulerOptions , CancellationToken cancellationToken )
155204 {
156205 var timestamp = Stopwatch . GetTimestamp ( ) ;
157206 schedulerOptions . Seed = GenerateSeed ( schedulerOptions . Seed ) ;
158207
159- //1. Create filenames
208+ // 1. Create filenames
160209 var random = await _fileService . CreateRandomName ( ) ;
161- var output = $ "Output-{ random } ";
162- var outputImage = $ "{ output } .png";
163- var outputBlueprint = $ "{ output } .json";
210+ var maskImage = $ "{ random } -mask.png";
211+ var inputImage = $ "{ random } -input.png";
212+ var outputImage = $ "{ random } -output.png";
213+ var outputBlueprint = $ "{ random } -output.json";
164214 var outputImageUrl = await _fileService . CreateOutputUrl ( outputImage ) ;
165215 var outputImageFile = await _fileService . UrlToPhysicalPath ( outputImageUrl ) ;
166-
167- //2. Generate blueprint
168216 var outputImageLink = await _fileService . CreateOutputUrl ( outputImage , false ) ;
169- var blueprint = new ImageBlueprint ( promptOptions , schedulerOptions , outputImageLink ) ;
217+ var inputImageLink = await _fileService . CreateOutputUrl ( inputImage , false ) ;
218+ var maskImageLink = await _fileService . CreateOutputUrl ( maskImage , false ) ;
219+
220+ // 2. Save Input Image to disk
221+ var inputImageFile = await _fileService . SaveImageFile ( promptOptions . InputImage . ImageBase64 , inputImage ) ;
222+ if ( inputImageFile is null )
223+ return new StableDiffusionResult ( "Failed to save input image" ) ;
224+
225+ // 3. Save Mask Image to disk
226+ var maskImageFile = await _fileService . SaveImageFile ( promptOptions . InputImageMask . ImageBase64 , maskImage ) ;
227+ if ( maskImageFile is null )
228+ return new StableDiffusionResult ( "Failed to save mask image" ) ;
229+
230+ // 4. Generate and save blueprint file
231+ var blueprint = new ImageBlueprint ( promptOptions , schedulerOptions , outputImageLink , inputImageLink , maskImageLink ) ;
170232 var bluprintFile = await _fileService . SaveBlueprintFile ( blueprint , outputBlueprint ) ;
171233 if ( bluprintFile is null )
172234 return new StableDiffusionResult ( "Failed to save blueprint" ) ;
173235
174- //3 . Run stable diffusion
236+ // 5 . Run stable diffusion
175237 if ( ! await RunStableDiffusion ( promptOptions , schedulerOptions , outputImageFile , cancellationToken ) )
176238 return new StableDiffusionResult ( "Failed to run stable diffusion" ) ;
177239
178- //4 . Return result
240+ // 6 . Return result
179241 return new StableDiffusionResult ( outputImage , outputImageUrl , blueprint , bluprintFile . Filename , bluprintFile . FileUrl , GetElapsed ( timestamp ) ) ;
180242 }
181243
0 commit comments