Skip to content

Commit a17d0e4

Browse files
kzukzu
committed
Improve image generation with multi-image editing, options, and usage
- Add GrokImageGenerationOptions subclass with AspectRatio and Resolution properties for strongly-typed control over grok-imagine model features - Fix GrokImageGenerator to use singular Image field for 1 reference image and the repeated Images field for 2+ reference images (multi-image editing) - Map ImageResponse.Usage to ImageGenerationResponse.Usage via UsageDetails - Add unit and integration tests for all three improvements
1 parent 6535199 commit a17d0e4

File tree

3 files changed

+244
-17
lines changed

3 files changed

+244
-17
lines changed

‎src/xAI.Tests/ImageGeneratorTests.cs‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
1+
using Grpc.Core;
12
using Microsoft.Extensions.AI;
3+
using Moq;
24
using Tests.Client.Helpers;
35
using xAI;
6+
using xAI.Protocol;
47
using static ConfigurationExtensions;
58

69
namespace xAI.Tests;
@@ -164,6 +167,177 @@ public async Task GenerateImage_ResponseContainsRawRepresentation()
164167
output.WriteLine($"Model used: {rawResponse.Model}");
165168
}
166169

170+
[SecretsFact("XAI_API_KEY")]
171+
public async Task GenerateImage_WithAspectRatioAndResolution_ReturnsImageContent()
172+
{
173+
var imageGenerator = new GrokClient(Configuration["XAI_API_KEY"]!)
174+
.AsIImageGenerator("grok-imagine-image");
175+
176+
var request = new ImageGenerationRequest("A cinematic skyline at sunrise");
177+
var options = new GrokImageGenerationOptions
178+
{
179+
ResponseFormat = ImageGenerationResponseFormat.Uri,
180+
AspectRatio = ImageAspectRatio.ImgAspectRatio169,
181+
Resolution = ImageResolution.ImgResolution1K,
182+
Count = 1
183+
};
184+
185+
var response = await imageGenerator.GenerateAsync(request, options);
186+
187+
Assert.NotNull(response);
188+
Assert.NotEmpty(response.Contents);
189+
Assert.Single(response.Contents);
190+
191+
var image = Assert.IsType<UriContent>(response.Contents.First());
192+
Assert.Equal("image/jpeg", image.MediaType);
193+
output.WriteLine($"Generated image URL: {image.Uri}");
194+
}
195+
196+
[LocalFact("XAI_API_KEY")]
197+
public async Task GenerateImage_WithMultiImageEdit_ReturnsImageContent()
198+
{
199+
var imageGenerator = new GrokClient(Configuration["XAI_API_KEY"]!)
200+
.AsIImageGenerator("grok-imagine-image");
201+
202+
var seedOptions = new ImageGenerationOptions
203+
{
204+
ResponseFormat = ImageGenerationResponseFormat.Uri,
205+
Count = 2
206+
};
207+
208+
var seeds = await imageGenerator.GenerateAsync(
209+
new ImageGenerationRequest("Two stylized character portraits"), seedOptions);
210+
211+
Assert.NotNull(seeds);
212+
Assert.NotEmpty(seeds.Contents);
213+
Assert.Equal(2, seeds.Contents.Count);
214+
215+
var originals = seeds.Contents.Take(2).ToList();
216+
var edit = await imageGenerator.GenerateAsync(
217+
new ImageGenerationRequest("Combine both portraits into a single movie poster", originals),
218+
new ImageGenerationOptions
219+
{
220+
ResponseFormat = ImageGenerationResponseFormat.Uri,
221+
Count = 1
222+
});
223+
224+
Assert.NotNull(edit);
225+
Assert.NotEmpty(edit.Contents);
226+
Assert.Single(edit.Contents);
227+
var image = Assert.IsType<UriContent>(edit.Contents.First());
228+
Assert.Equal("image/jpeg", image.MediaType);
229+
output.WriteLine($"Edited image URL: {image.Uri}");
230+
}
231+
232+
[Fact]
233+
public async Task GenerateImage_WithOneOriginalImage_SetsImageField()
234+
{
235+
GenerateImageRequest? capturedRequest = null;
236+
var client = new Mock<Image.ImageClient>(MockBehavior.Strict);
237+
client.Setup(x => x.GenerateImageAsync(It.IsAny<GenerateImageRequest>(), null, null, CancellationToken.None))
238+
.Callback<GenerateImageRequest, Metadata?, DateTime?, CancellationToken>((req, _, _, _) => capturedRequest = req)
239+
.Returns(CallHelpers.CreateAsyncUnaryCall(new ImageResponse
240+
{
241+
Images =
242+
{
243+
new GeneratedImage { Url = "https://example.com/generated.jpg" }
244+
}
245+
}));
246+
247+
var imageGenerator = client.Object.AsIImageGenerator("grok-imagine-image");
248+
var source = new UriContent(new Uri("https://example.com/source.jpg"), "image/jpeg");
249+
250+
await imageGenerator.GenerateAsync(new ImageGenerationRequest("Edit this image", [source]));
251+
252+
Assert.NotNull(capturedRequest);
253+
Assert.NotNull(capturedRequest.Image);
254+
Assert.Equal("https://example.com/source.jpg", capturedRequest.Image.ImageUrl);
255+
Assert.Empty(capturedRequest.Images);
256+
}
257+
258+
[Fact]
259+
public async Task GenerateImage_WithMultipleOriginalImages_SetsImagesField()
260+
{
261+
GenerateImageRequest? capturedRequest = null;
262+
var client = new Mock<Image.ImageClient>(MockBehavior.Strict);
263+
client.Setup(x => x.GenerateImageAsync(It.IsAny<GenerateImageRequest>(), null, null, CancellationToken.None))
264+
.Callback<GenerateImageRequest, Metadata?, DateTime?, CancellationToken>((req, _, _, _) => capturedRequest = req)
265+
.Returns(CallHelpers.CreateAsyncUnaryCall(new ImageResponse
266+
{
267+
Images =
268+
{
269+
new GeneratedImage { Url = "https://example.com/generated.jpg" }
270+
}
271+
}));
272+
273+
var imageGenerator = client.Object.AsIImageGenerator("grok-imagine-image");
274+
var first = new UriContent(new Uri("https://example.com/source-1.jpg"), "image/jpeg");
275+
var second = new UriContent(new Uri("https://example.com/source-2.jpg"), "image/jpeg");
276+
277+
await imageGenerator.GenerateAsync(new ImageGenerationRequest("Blend these images", [first, second]));
278+
279+
Assert.NotNull(capturedRequest);
280+
Assert.Null(capturedRequest.Image);
281+
Assert.Equal(2, capturedRequest.Images.Count);
282+
Assert.Equal("https://example.com/source-1.jpg", capturedRequest.Images[0].ImageUrl);
283+
Assert.Equal("https://example.com/source-2.jpg", capturedRequest.Images[1].ImageUrl);
284+
}
285+
286+
[Fact]
287+
public async Task GenerateImage_WithAspectRatioOption_SetsProtocolAspectRatio()
288+
{
289+
GenerateImageRequest? capturedRequest = null;
290+
var client = new Mock<Image.ImageClient>(MockBehavior.Strict);
291+
client.Setup(x => x.GenerateImageAsync(It.IsAny<GenerateImageRequest>(), null, null, CancellationToken.None))
292+
.Callback<GenerateImageRequest, Metadata?, DateTime?, CancellationToken>((req, _, _, _) => capturedRequest = req)
293+
.Returns(CallHelpers.CreateAsyncUnaryCall(new ImageResponse
294+
{
295+
Images =
296+
{
297+
new GeneratedImage { Url = "https://example.com/generated.jpg" }
298+
}
299+
}));
300+
301+
var imageGenerator = client.Object.AsIImageGenerator("grok-imagine-image");
302+
303+
await imageGenerator.GenerateAsync(
304+
new ImageGenerationRequest("Wide composition"),
305+
new GrokImageGenerationOptions { AspectRatio = ImageAspectRatio.ImgAspectRatio169 });
306+
307+
Assert.NotNull(capturedRequest);
308+
Assert.True(capturedRequest.HasAspectRatio);
309+
Assert.Equal(ImageAspectRatio.ImgAspectRatio169, capturedRequest.AspectRatio);
310+
}
311+
312+
[Fact]
313+
public async Task GenerateImage_MapsProtocolUsageToResponseUsage()
314+
{
315+
var client = new Mock<Image.ImageClient>(MockBehavior.Strict);
316+
client.Setup(x => x.GenerateImageAsync(It.IsAny<GenerateImageRequest>(), null, null, CancellationToken.None))
317+
.Returns(CallHelpers.CreateAsyncUnaryCall(new ImageResponse
318+
{
319+
Images =
320+
{
321+
new GeneratedImage { Url = "https://example.com/generated.jpg" }
322+
},
323+
Usage = new SamplingUsage
324+
{
325+
PromptTokens = 11,
326+
CompletionTokens = 7,
327+
TotalTokens = 18
328+
}
329+
}));
330+
331+
var imageGenerator = client.Object.AsIImageGenerator("grok-imagine-image");
332+
var response = await imageGenerator.GenerateAsync(new ImageGenerationRequest("Test usage mapping"));
333+
334+
Assert.NotNull(response);
335+
Assert.NotNull(response.Usage);
336+
Assert.Equal(11, response.Usage.InputTokenCount);
337+
Assert.Equal(7, response.Usage.OutputTokenCount);
338+
Assert.Equal(18, response.Usage.TotalTokenCount);
339+
}
340+
167341
[Fact]
168342
public async Task GenerateImage_WithNullRequest_ThrowsArgumentNullException()
169343
{

‎src/xAI/GrokImageGenerationOptions.cs‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
using Microsoft.Extensions.AI;
2+
using xAI.Protocol;
3+
4+
namespace xAI;
5+
6+
/// <summary>Grok-specific image generation options that extend the base <see cref="ImageGenerationOptions"/>.</summary>
7+
/// <remarks>
8+
/// These options map to image.proto fields and are only supported by grok-imagine models.
9+
/// If not specified, the API defaults to 1:1 aspect ratio and 1k resolution.
10+
/// </remarks>
11+
public class GrokImageGenerationOptions : ImageGenerationOptions
12+
{
13+
/// <summary>Optional aspect ratio for image generation and editing.</summary>
14+
/// <remarks>
15+
/// Proto default is 1:1 when this option is not specified.
16+
/// Auto aspect ratio is only supported for image generation with a thinking upsampler.
17+
/// This option is only supported by grok-imagine models.
18+
/// </remarks>
19+
public ImageAspectRatio? AspectRatio { get; set; }
20+
21+
/// <summary>Optional resolution for image generation and editing.</summary>
22+
/// <remarks>
23+
/// Proto default is 1k when this option is not specified.
24+
/// 2k output is generated at 1k and then upscaled with super-resolution.
25+
/// This option is only supported by grok-imagine models.
26+
/// </remarks>
27+
public ImageResolution? Resolution { get; set; }
28+
}

‎src/xAI/GrokImageGenerator.cs‎

Lines changed: 42 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -69,29 +69,27 @@ public async Task<ImageGenerationResponse> GenerateAsync(
6969
_ => throw new ArgumentException($"Unsupported response format: {options?.ResponseFormat}", nameof(options))
7070
};
7171

72+
if (options is GrokImageGenerationOptions grokOptions)
73+
{
74+
if (grokOptions.AspectRatio is { } aspectRatio) protocolRequest.AspectRatio = aspectRatio;
75+
if (grokOptions.Resolution is { } resolution) protocolRequest.Resolution = resolution;
76+
}
77+
7278
// Handle image editing if original images are provided
73-
if (request.OriginalImages?.FirstOrDefault() is { } originalImage)
79+
if (request.OriginalImages?.ToList() is { Count: > 0 } originalImages)
7480
{
75-
if (originalImage is DataContent dataContent)
81+
if (originalImages.Count == 1)
7682
{
77-
var imageUrl = dataContent.Uri?.ToString();
78-
if (imageUrl == null && dataContent.Data.Length > 0)
79-
imageUrl = $"data:{dataContent.MediaType ?? DefaultInputContentType};base64,{Convert.ToBase64String(dataContent.Data.ToArray())}";
80-
81-
if (imageUrl != null)
82-
{
83-
protocolRequest.Image = new ImageUrlContent
84-
{
85-
ImageUrl = imageUrl
86-
};
87-
}
83+
if (MapToImageUrlContent(originalImages[0]) is { } image)
84+
protocolRequest.Image = image;
8885
}
89-
else if (originalImage is UriContent uriContent)
86+
else
9087
{
91-
protocolRequest.Image = new ImageUrlContent
88+
foreach (var originalImage in originalImages)
9289
{
93-
ImageUrl = uriContent.Uri.ToString()
94-
};
90+
if (MapToImageUrlContent(originalImage) is { } image)
91+
protocolRequest.Images.Add(image);
92+
}
9593
}
9694
}
9795

@@ -156,6 +154,33 @@ static ImageGenerationResponse ToImageGenerationResponse(ImageResponse response)
156154
return new ImageGenerationResponse(contents)
157155
{
158156
RawRepresentation = response,
157+
Usage = MapToUsage(response.Usage),
158+
};
159+
}
160+
161+
static ImageUrlContent? MapToImageUrlContent(AIContent content) => content switch
162+
{
163+
DataContent dataContent => MapToImageUrlContent(dataContent),
164+
UriContent uriContent => new ImageUrlContent { ImageUrl = uriContent.Uri.ToString() },
165+
_ => throw new ArgumentException($"Unsupported original image content type: {content.GetType()}", nameof(content)),
166+
};
167+
168+
static ImageUrlContent? MapToImageUrlContent(DataContent dataContent)
169+
{
170+
var imageUrl = dataContent.Uri?.ToString();
171+
if (imageUrl == null && dataContent.Data.Length > 0)
172+
imageUrl = $"data:{dataContent.MediaType ?? DefaultInputContentType};base64,{Convert.ToBase64String(dataContent.Data.ToArray())}";
173+
174+
return imageUrl == null ? null : new ImageUrlContent
175+
{
176+
ImageUrl = imageUrl
159177
};
160178
}
179+
180+
static UsageDetails? MapToUsage(SamplingUsage usage) => usage == null ? null : new()
181+
{
182+
InputTokenCount = usage.PromptTokens,
183+
OutputTokenCount = usage.CompletionTokens,
184+
TotalTokenCount = usage.TotalTokens
185+
};
161186
}

0 commit comments

Comments
 (0)