Add image-generation support and UI

Madzionator · Madzionator · commit 07034d12d3ff · 2026-02-26T14:50:16.000+01:00
Introduce image-generation capability across the app: add IImageGenerationModel and HasImageGeneration on AIModel; mark cloud models (DALL·E3, new gpt-image-1 and grok-2-image) as image generators. Update UI to render generated images with download and copy-to-clipboard actions (Home.razor changes, CopyImageToClipboard interop + editor.js). Improve visual/model detection in Utils to use ModelRegistry with a fallback set of known image-generation IDs. Increase SignalR hub max message size to 10MB (Program.cs) to allow larger image transfers and add CSS for generated image layout and controls.
diff --git a/src/MaIN.Domain/Models/Abstract/AIModel.cs b/src/MaIN.Domain/Models/Abstract/AIModel.cs
@@ -34,6 +34,9 @@ public abstract record AIModel(
     
     /// <summary> Checks if model supports vision/image input. </summary>
     public bool HasVision => this is IVisionModel;
+
+    /// <summary> Checks if model generates images from text prompts. </summary>
+    public bool HasImageGeneration => this is IImageGenerationModel;
 }
 
 /// <summary> Base class for local models. </summary>
diff --git a/src/MaIN.Domain/Models/Abstract/IModelCapabilities.cs b/src/MaIN.Domain/Models/Abstract/IModelCapabilities.cs
@@ -45,3 +45,8 @@ public interface IEmbeddingModel
 /// Interface for models that support text-to-speech.
 /// </summary>
 public interface ITTSModel;
+
+/// <summary>
+/// Interface for models that generate images from text prompts.
+/// </summary>
+public interface IImageGenerationModel;
diff --git a/src/MaIN.Domain/Models/Concrete/CloudModels.cs b/src/MaIN.Domain/Models/Concrete/CloudModels.cs
@@ -31,7 +31,14 @@ public sealed record DallE3() : CloudModel(
     BackendType.OpenAi,
     "DALL-E 3",
     4000,
-    "Advanced image generation model from OpenAI");
+    "Advanced image generation model from OpenAI"), IImageGenerationModel;
+
+public sealed record GptImage1() : CloudModel(
+    "gpt-image-1",
+    BackendType.OpenAi,
+    "GPT Image 1",
+    4000,
+    "OpenAI's latest image generation model"), IImageGenerationModel;
 
 // ===== Anthropic Models =====
 
@@ -74,6 +81,13 @@ public sealed record Grok3Beta() : CloudModel(
     ModelDefaults.DefaultMaxContextWindow,
     "xAI latest Grok model in beta testing phase");
 
+public sealed record GrokImage() : CloudModel(
+    "grok-2-image",
+    BackendType.Xai,
+    "Grok 2 Image",
+    4000,
+    "xAI image generation model"), IImageGenerationModel;
+
 // ===== GroqCloud Models =====
 
 public sealed record Llama3_1_8bInstant() : CloudModel(
diff --git a/src/MaIN.InferPage/Components/Layout/NavBar.razor b/src/MaIN.InferPage/Components/Layout/NavBar.razor
@@ -29,6 +29,15 @@
             Color="#000"
             Style="margin-left: 10px">Reasoning ✨</FluentBadge>
     }
+    @if (Utils.Visual)
+    {
+        <FluentBadge
+            Appearance="Appearance.Neutral"
+            Fill="highlight"
+            BackgroundColor="#7c3aed"
+            Color="#fff"
+            Style="margin-left: 10px">Visual 🎨</FluentBadge>
+    }
     <div  style="margin-left: auto; align-self: flex-end;">
         <FluentButton Style="padding: 10px; background-color: transparent;"
                       BackgroundColor="rgba(0, 0, 0, 0)"
diff --git a/src/MaIN.InferPage/Components/Pages/Home.razor b/src/MaIN.InferPage/Components/Pages/Home.razor
@@ -36,37 +36,10 @@
         {
             @if (conversation.Message.Role != "System")
             {
-                @if (Chat.Visual)
-                {
-                    <FluentBadge Class="@(conversation.Message.Role == "User" ? "message-role-user" : "message-role-bot")"
-                                 Appearance="Appearance.Accent">
-                        @(conversation.Message.Role == "User" ? "User" : Utils.Model)
-                    </FluentBadge>
+                <FluentCard class="@(conversation.Message.Role == "User" ? "message-card user-message" : "message-card bot-message")">
                     @if (conversation.Message.Role == "User")
                     {
-                        <FluentCard class="message-card user-message">
-                            @conversation.Message.Content
-                        </FluentCard>
-                    }
-                    else
-                    {
-                        <FluentCard class="message-card-img bot-message"
-                                    Style="height: 30rem !important; width: 30rem !important; ">
-                            <div>
-                                <a href="data:image/png;base64,@Convert.ToBase64String(conversation.Message.Image!)"
-                                   style="cursor: -webkit-zoom-in; cursor: zoom-in;" target="_blank">
-                                    <img src="data:image/png;base64,@Convert.ToBase64String(conversation.Message.Image!)"
-                                         style="object-fit: fill; width:100%; height:100%;"
-                                         alt="imageResponse"/>
-                                </a>
-                            </div>
-                        </FluentCard>
-                    }
-                }
-                else
-                {
-                    <FluentCard class="@(conversation.Message.Role == "User" ? "message-card user-message" : "message-card bot-message")">
-                        @if (conversation.Message.Role == "User" && (conversation.AttachedFiles.Any() || conversation.AttachedImages.Any()))
+                        @if (conversation.AttachedFiles.Any() || conversation.AttachedImages.Any())
                         {
                             <div class="attached-files-display">
                                 @foreach (var image in conversation.AttachedImages)
@@ -84,10 +57,36 @@
                                 }
                             </div>
                         }
-                        @if (conversation.Message.Role == "User")
+                        <div>
+                            @((MarkupString)Markdown.ToHtml(conversation.Message.Content ?? string.Empty, _markdownPipeline))
+                        </div>
+                    }
+                    else
+                    {
+                        @if (conversation.Message.Images?.Any() == true)
                         {
-                            <div>
-                                @((MarkupString)Markdown.ToHtml(conversation.Message.Content ?? string.Empty, _markdownPipeline))
+                            <div class="generated-images">
+                                @foreach (var imageBytes in conversation.Message.Images)
+                                {
+                                    var b64 = Convert.ToBase64String(imageBytes);
+                                    <div class="image-wrapper">
+                                        <a href="data:image/png;base64,@b64" target="_blank">
+                                            <img src="data:image/png;base64,@b64" class="generated-image" alt="generated image" />
+                                        </a>
+                                        <div class="image-actions">
+                                            <a href="data:image/png;base64,@b64" download="generated-image.png"
+                                               class="image-action-btn" title="Download">
+                                                <FluentIcon Value="@(new Icons.Filled.Size24.ArrowDownload())"
+                                                            Style="fill: var(--accent-base-color);" />
+                                            </a>
+                                            <span class="image-action-btn" title="Copy to clipboard"
+                                                  @onclick="@(async () => await CopyImageToClipboard(b64))">
+                                                <FluentIcon Value="@(new Icons.Filled.Size24.Copy())"
+                                                            Style="fill: var(--accent-base-color);" />
+                                            </span>
+                                        </div>
+                                    </div>
+                                }
                             </div>
                         }
                         else
@@ -114,33 +113,30 @@
                                 @((MarkupString)Markdown.ToHtml(GetMessageContent(conversation.Message), _markdownPipeline))
                             </div>
                         }
-                    </FluentCard>
-                }
+                    }
+                </FluentCard>
             }
         }
         @if (_isLoading)
         {
-            @if (Chat.Visual)
+            @if (Utils.Visual)
             {
                 <span class="message-role-bot" style="font-style: italic; font-size: small">This might take a while...</span>
             }
-            else
+            else if (_incomingMessage != null || _incomingReasoning != null)
             {
-                @if (_incomingMessage != null || _incomingReasoning != null)
-                {
-                    <FluentCard class="message-card bot-message">
-                        @if (_isThinking)
-                        {
-                            <span class="thinker">
-                                @((MarkupString)Markdown.ToHtml(_incomingReasoning ?? string.Empty, _markdownPipeline))
-                            </span>
-                        }
-                        else
-                        {
-                            @((MarkupString)Markdown.ToHtml(_incomingMessage ?? string.Empty, _markdownPipeline))
-                        }
-                    </FluentCard>
-                }
+                <FluentCard class="message-card bot-message">
+                    @if (_isThinking)
+                    {
+                        <span class="thinker">
+                            @((MarkupString)Markdown.ToHtml(_incomingReasoning ?? string.Empty, _markdownPipeline))
+                        </span>
+                    }
+                    else
+                    {
+                        @((MarkupString)Markdown.ToHtml(_incomingMessage ?? string.Empty, _markdownPipeline))
+                    }
+                </FluentCard>
             }
         }
         <div id="bottom" @ref="_bottomElement"></div>
@@ -277,7 +273,7 @@
             }
 
             ctx = Utils.Visual
-                ? AIHub.Chat().EnableVisual()
+                ? AIHub.Chat().WithModel(model).EnableVisual()
                 : AIHub.Chat().WithModel(model);
         }
         catch (MaINCustomException ex)
@@ -628,6 +624,18 @@
         }).ToList();
     }
 
+    private async Task CopyImageToClipboard(string base64)
+    {
+        try
+        {
+            await JS.InvokeVoidAsync("editorManager.copyImageToClipboard", base64);
+        }
+        catch
+        {
+            // silently ignore — clipboard API may not be available in all browsers
+        }
+    }
+
     public void Dispose()
     {
         _cancellationTokenSource?.Dispose();
diff --git a/src/MaIN.InferPage/Program.cs b/src/MaIN.InferPage/Program.cs
@@ -8,7 +8,11 @@
 
 var builder = WebApplication.CreateBuilder(args);
 builder.Services.AddRazorComponents()
-    .AddInteractiveServerComponents();
+    .AddInteractiveServerComponents()
+    .AddHubOptions(options =>
+    {
+        options.MaximumReceiveMessageSize = 10 * 1024 * 1024; // 10 MB
+    });
 builder.Services.AddFluentUIComponents();
 
 try
diff --git a/src/MaIN.InferPage/Utils.cs b/src/MaIN.InferPage/Utils.cs
@@ -1,18 +1,36 @@
 using MaIN.Domain.Configuration;
 using MaIN.Domain.Entities;
+using MaIN.Domain.Models.Abstract;
 
 namespace MaIN.InferPage;
 
 public static class Utils
 {
     public static BackendType BackendType { get; set; } = BackendType.Self;
     public static bool HasApiKey { get; set; }
+    public static string? Path { get; set; }
     public static bool IsLocal => BackendType == BackendType.Self || (BackendType == BackendType.Ollama && !HasApiKey);
     public static string? Model = "gemma3-4b";
     public static bool Reason { get; set; }
-    public static bool Visual => VisualModels.Contains(Model);
-    private static readonly string[] VisualModels = ["FLUX.1_Shnell", "FLUX.1", "dall-e-3", "dall-e", "imagen", "imagen-3"]; //user might type different names
-    public static string? Path { get; set; }
+    public static bool Visual
+    {
+        get
+        {
+            if (string.IsNullOrEmpty(Model)) return false;
+            if (ModelRegistry.TryGetById(Model, out var m))
+                return m is IImageGenerationModel;
+            return ImageGenerationModels.Contains(Model); // fallback for unregistered models (e.g. FLUX via separate server)
+        }
+    }
+
+    private static readonly HashSet<string> ImageGenerationModels =
+    [
+        "FLUX.1_Shnell", "FLUX.1",
+        "dall-e-3", "dall-e",
+        "gpt-image-1",
+        "imagen", "imagen-3", "imagen-4", "imagen-4-fast",
+        "grok-2-image"
+    ];
 }
 
 public class MessageExt
diff --git a/src/MaIN.InferPage/wwwroot/editor.js b/src/MaIN.InferPage/wwwroot/editor.js
@@ -109,5 +109,10 @@ window.editorManager = {
         } catch (err) {
             try { await dotNetHelper.invokeMethodAsync('OnDragLeave'); } catch {}
         }
+    },
+    copyImageToClipboard: async (base64) => {
+        const res = await fetch(`data:image/png;base64,${base64}`);
+        const blob = await res.blob();
+        await navigator.clipboard.write([new ClipboardItem({ 'image/png': blob })]);
     }
 };
diff --git a/src/MaIN.InferPage/wwwroot/home.css b/src/MaIN.InferPage/wwwroot/home.css
@@ -406,3 +406,58 @@ body {
     30%, 50%, 70% { transform: translate3d(-4px, 0, 0); }
     40%, 60% { transform: translate3d(4px, 0, 0); }
 }
+
+/* Generated image display */
+.generated-images {
+    display: flex;
+    flex-direction: column;
+    gap: 12px;
+    width: 100%;
+}
+
+.image-wrapper {
+    position: relative;
+    width: 100%;
+}
+
+/* Fix: <a> is inline by default — make it block so width: 100% on <img> works correctly */
+.image-wrapper > a {
+    display: block;
+    width: 100%;
+}
+
+.generated-image {
+    width: 100%;
+    height: auto;
+    border-radius: 8px;
+    display: block;
+    cursor: zoom-in;
+}
+
+/* Action buttons overlaid on bottom-left corner of the image */
+.image-actions {
+    display: flex;
+    gap: 6px;
+    position: absolute;
+    bottom: 10px;
+    left: 10px;
+}
+
+.image-action-btn {
+    display: inline-flex;
+    align-items: center;
+    justify-content: center;
+    padding: 8px;
+    border-radius: 8px;
+    opacity: 0.85;
+    transition: opacity 0.15s, background-color 0.15s;
+    color: var(--accent-base-color);
+    text-decoration: none;
+    cursor: pointer;
+    background-color: rgba(0, 0, 0, 0.4);
+}
+
+.image-action-btn:hover {
+    opacity: 1;
+    background-color: rgba(0, 0, 0, 0.6);
+}

Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,9 @@ public abstract record AIModel(`
`34`	`34`
`35`	`35`	`/// <summary> Checks if model supports vision/image input. </summary>`
`36`	`36`	`public bool HasVision => this is IVisionModel;`
	`37`	`+`
	`38`	`+ /// <summary> Checks if model generates images from text prompts. </summary>`
	`39`	`+ public bool HasImageGeneration => this is IImageGenerationModel;`
`37`	`40`	`}`
`38`	`41`
`39`	`42`	`/// <summary> Base class for local models. </summary>`
Original file line number	Diff line number	Diff line change
`@@ -109,5 +109,10 @@ window.editorManager = {`
`109`	`109`	`} catch (err) {`
`110`	`110`	`try { await dotNetHelper.invokeMethodAsync('OnDragLeave'); } catch {}`
`111`	`111`	`}`
	`112`	`+ },`
	`113`	`+ copyImageToClipboard: async (base64) => {`
	`114`	+ const res = await fetch(`data:image/png;base64,${base64}`);
	`115`	`+ const blob = await res.blob();`
	`116`	`+ await navigator.clipboard.write([new ClipboardItem({ 'image/png': blob })]);`
`112`	`117`	`}`
`113`	`118`	`};`