Skip to content

Commit cf88b55

Browse files
committed
Function: Add tensorBlock Image Generation Support
1 parent 1b1efa2 commit cf88b55

File tree

3 files changed

+137
-35
lines changed

3 files changed

+137
-35
lines changed

src/components/pages/ImageGenerationPage.tsx

Lines changed: 92 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { ChevronDown, RefreshCw, Settings, Zap } from "lucide-react";
77
import { useTranslation } from "react-i18next";
88
import { AIService } from "../../services/ai-service";
99
import { OPENAI_PROVIDER_NAME } from "../../services/providers/openai-service";
10+
import { FORGE_PROVIDER_NAME as TENSORBLOCK_PROVIDER_NAME } from "../../services/providers/forge-service";
1011
import { ImageGenerationManager, ImageGenerationStatus, ImageGenerationHandler } from "../../services/image-generation-handler";
1112
import { DatabaseIntegrationService } from "../../services/database-integration";
1213
import { ImageGenerationResult } from "../../types/image";
@@ -29,9 +30,12 @@ export const ImageGenerationPage = () => {
2930
const [historyResults, setHistoryResults] = useState<ImageGenerationResult[]>([]);
3031
const [isLoadingHistory, setIsLoadingHistory] = useState(true);
3132
const [isSettingsOpen, setIsSettingsOpen] = useState(false);
33+
const [selectedProvider, setSelectedProvider] = useState(OPENAI_PROVIDER_NAME);
3234

3335
const settingsButtonRef = useRef<HTMLButtonElement>(null);
3436
const settingsPopupRef = useRef<HTMLDivElement>(null);
37+
const providerDropdownRef = useRef<HTMLDivElement>(null);
38+
const [isProviderDropdownOpen, setIsProviderDropdownOpen] = useState(false);
3539

3640
// Load image generation history from database
3741
const refreshImageHistory = useCallback(async () => {
@@ -50,7 +54,7 @@ export const ImageGenerationPage = () => {
5054
}
5155
}, []);
5256

53-
// Initialize image generation manager
57+
// Initialize image generation manager and load settings
5458
useEffect(() => {
5559
const initialize = async () => {
5660
// Initialize image generation manager
@@ -67,6 +71,16 @@ export const ImageGenerationPage = () => {
6771
const dbService = DatabaseIntegrationService.getInstance();
6872
await dbService.initialize();
6973

74+
// Load settings
75+
const settingsService = SettingsService.getInstance();
76+
await settingsService.initialize();
77+
78+
// Load saved provider preference
79+
const settings = settingsService.getSettings();
80+
if (settings.imageGenerationProvider) {
81+
setSelectedProvider(settings.imageGenerationProvider);
82+
}
83+
7084
// Load image generation history from database
7185
await refreshImageHistory();
7286

@@ -82,18 +96,24 @@ export const ImageGenerationPage = () => {
8296

8397
// Check if API key is available
8498
useEffect(() => {
85-
setIsApiKeyMissing(!SettingsService.getInstance().getApiKey());
99+
const checkApiKey = () => {
100+
// Check if the selected provider has an API key
101+
const hasApiKey = !!SettingsService.getInstance().getApiKey(selectedProvider);
102+
setIsApiKeyMissing(!hasApiKey);
103+
};
104+
105+
checkApiKey();
86106

87107
const handleSettingsChange = () => {
88-
setIsApiKeyMissing(!SettingsService.getInstance().getApiKey());
108+
checkApiKey();
89109
};
90110

91111
window.addEventListener(SETTINGS_CHANGE_EVENT, handleSettingsChange);
92112

93113
return () => {
94114
window.removeEventListener(SETTINGS_CHANGE_EVENT, handleSettingsChange);
95115
};
96-
}, []);
116+
}, [selectedProvider]);
97117

98118
// Handle clicks outside the settings popup
99119
useEffect(() => {
@@ -105,6 +125,14 @@ export const ImageGenerationPage = () => {
105125
!settingsButtonRef.current.contains(event.target as Node)
106126
) {
107127
setIsSettingsOpen(false);
128+
setIsProviderDropdownOpen(false);
129+
} else if (
130+
providerDropdownRef.current &&
131+
!providerDropdownRef.current.contains(event.target as Node) &&
132+
event.target instanceof Element &&
133+
!event.target.closest('.provider-dropdown-toggle')
134+
) {
135+
setIsProviderDropdownOpen(false);
108136
}
109137
};
110138

@@ -114,18 +142,24 @@ export const ImageGenerationPage = () => {
114142
};
115143
}, []);
116144

117-
// Handle generating an image using OpenAI's DALL-E 3
145+
// Handle generating an image using selected provider
118146
const handleGenerateImage = async () => {
119147
if (!prompt.trim()) return;
120148

121149
setError(null);
122150

123151
try {
124-
// Get the OpenAI service from AIService
125-
const openaiService = AIService.getInstance().getProvider(OPENAI_PROVIDER_NAME);
152+
let providerService;
153+
154+
// Get the appropriate service based on selected provider
155+
if (selectedProvider === TENSORBLOCK_PROVIDER_NAME) {
156+
providerService = AIService.getInstance().getProvider(TENSORBLOCK_PROVIDER_NAME);
157+
} else {
158+
providerService = AIService.getInstance().getProvider(OPENAI_PROVIDER_NAME);
159+
}
126160

127-
if (!openaiService) {
128-
throw new Error("OpenAI service not available");
161+
if (!providerService) {
162+
throw new Error(`${selectedProvider} service not available`);
129163
}
130164

131165
// Create a new generation handler
@@ -135,7 +169,7 @@ export const ImageGenerationPage = () => {
135169
seed: randomSeed,
136170
number: imageCount,
137171
aspectRatio: aspectRatio,
138-
provider: OPENAI_PROVIDER_NAME,
172+
provider: selectedProvider,
139173
model: "dall-e-3",
140174
});
141175

@@ -153,7 +187,7 @@ export const ImageGenerationPage = () => {
153187
};
154188

155189
// Generate the image
156-
const images = await openaiService.getImageGeneration(prompt, {
190+
const images = await providerService.getImageGeneration(prompt, {
157191
size: sizeMap[aspectRatio] || "1024x1024",
158192
aspectRatio: aspectRatio as `${number}:${number}`,
159193
style: "vivid"
@@ -223,6 +257,24 @@ export const ImageGenerationPage = () => {
223257
// Toggle settings popup
224258
const toggleSettings = () => {
225259
setIsSettingsOpen(!isSettingsOpen);
260+
setIsProviderDropdownOpen(false);
261+
};
262+
263+
// Toggle provider dropdown
264+
const toggleProviderDropdown = () => {
265+
setIsProviderDropdownOpen(!isProviderDropdownOpen);
266+
};
267+
268+
// Handle provider selection
269+
const handleProviderSelect = async (provider: string) => {
270+
setSelectedProvider(provider);
271+
setIsProviderDropdownOpen(false);
272+
273+
// Save provider preference to settings
274+
const settingsService = SettingsService.getInstance();
275+
await settingsService.updateSettings({
276+
imageGenerationProvider: provider
277+
});
226278
};
227279

228280
return (
@@ -441,12 +493,38 @@ export const ImageGenerationPage = () => {
441493
</label>
442494
<div className="relative">
443495
<button
444-
className="flex items-center justify-between w-full p-3 text-left input-box"
445-
disabled={true}
496+
className="flex items-center justify-between w-full p-3 text-left provider-dropdown-toggle input-box"
497+
onClick={toggleProviderDropdown}
446498
>
447-
<span>OpenAI</span>
499+
<span>{selectedProvider}</span>
448500
<ChevronDown size={18} className="text-gray-500" />
449501
</button>
502+
503+
{isProviderDropdownOpen && (
504+
<div
505+
ref={providerDropdownRef}
506+
className="absolute z-20 w-full mt-1 bg-white border border-gray-300 rounded-md shadow-lg"
507+
>
508+
<ul className="py-1">
509+
<li
510+
className={`px-3 py-2 cursor-pointer hover:bg-gray-100 ${
511+
selectedProvider === OPENAI_PROVIDER_NAME ? 'bg-gray-50 font-medium' : ''
512+
}`}
513+
onClick={() => handleProviderSelect(OPENAI_PROVIDER_NAME)}
514+
>
515+
OpenAI
516+
</li>
517+
<li
518+
className={`px-3 py-2 cursor-pointer hover:bg-gray-100 ${
519+
selectedProvider === TENSORBLOCK_PROVIDER_NAME ? 'bg-gray-50 font-medium' : ''
520+
}`}
521+
onClick={() => handleProviderSelect(TENSORBLOCK_PROVIDER_NAME)}
522+
>
523+
TensorBlock
524+
</li>
525+
</ul>
526+
</div>
527+
)}
450528
</div>
451529
</div>
452530

src/services/providers/forge-service.ts

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -118,13 +118,29 @@ export class ForgeService implements AiServiceProvider {
118118
*/
119119
public async getImageGeneration(
120120
prompt: string,
121-
// eslint-disable-next-line @typescript-eslint/no-unused-vars
122121
options: {
123122
size?: `${number}x${number}`;
123+
aspectRatio?: `${number}:${number}`;
124124
style?: string;
125125
quality?: string;
126-
} = {}
127-
): Promise<string[]> {
128-
throw new Error('Not implemented');
126+
}
127+
): Promise<string[] | Uint8Array<ArrayBufferLike>[]> {
128+
129+
const imageModel = this.commonProviderHelper.ProviderInstance.imageModel('dall-e-3');
130+
131+
const result = await imageModel.doGenerate({
132+
prompt: prompt,
133+
n: 1,
134+
size: options.size || '1024x1024',
135+
aspectRatio: options.aspectRatio || '1:1',
136+
seed: 42,
137+
providerOptions: {
138+
"openai": {
139+
"style": options.style || 'vivid'
140+
}
141+
}
142+
});
143+
144+
return result.images;
129145
}
130146
}

src/services/settings-service.ts

Lines changed: 25 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -9,23 +9,31 @@ import { v4 as uuidv4 } from 'uuid';
99
*/
1010
const DEFAULT_SETTINGS: UserSettings = {
1111
providers: {
12-
// ['TensorBlock']: {
13-
// providerId: 'TensorBlock',
14-
// providerName: 'TensorBlock',
15-
// apiKey: '',
16-
// baseUrl: 'http://54.177.123.202:8000/v1',
17-
// customProvider: false,
18-
// models:[
19-
// {
20-
// modelName: 'GPT-4o',
21-
// modelId: 'gpt-4o',
22-
// modelCategory: 'GPT 4',
23-
// modelDescription: 'GPT-4o is the latest and most powerful model from OpenAI.',
24-
// modelCapabilities: [AIServiceCapability.TextCompletion, AIServiceCapability.WebSearch],
25-
// modelRefUUID: uuidv4(),
26-
// },
27-
// ]
28-
// },
12+
['TensorBlock']: {
13+
providerId: 'TensorBlock',
14+
providerName: 'TensorBlock',
15+
apiKey: '',
16+
baseUrl: 'http://54.177.123.202:8000/v1',
17+
customProvider: false,
18+
models:[
19+
{
20+
modelName: 'GPT-4o',
21+
modelId: 'gpt-4o',
22+
modelCategory: 'GPT 4',
23+
modelDescription: 'GPT-4o is the latest and most powerful model from OpenAI.',
24+
modelCapabilities: [AIServiceCapability.TextCompletion, AIServiceCapability.WebSearch],
25+
modelRefUUID: uuidv4(),
26+
},
27+
{
28+
modelName: 'DALL-E 3',
29+
modelId: 'dall-e-3',
30+
modelCategory: 'Image Generation',
31+
modelDescription: 'DALL-E 3 is OpenAI\'s advanced image generation model.',
32+
modelCapabilities: [AIServiceCapability.ImageGeneration],
33+
modelRefUUID: uuidv4(),
34+
},
35+
]
36+
},
2937
['OpenAI']: {
3038
providerId: 'OpenAI',
3139
providerName: 'OpenAI',

0 commit comments

Comments
 (0)