352 lines
8.9 KiB
Markdown
352 lines
8.9 KiB
Markdown
---
|
|
id: task-014
|
|
title: Implement WebGPU-based local image generation to reduce RunPod costs
|
|
status: To Do
|
|
assignee: []
|
|
created_date: '2025-12-04 11:46'
|
|
updated_date: '2025-12-04 11:47'
|
|
labels:
|
|
- performance
|
|
- cost-optimization
|
|
- webgpu
|
|
- ai
|
|
- image-generation
|
|
dependencies: []
|
|
priority: high
|
|
---
|
|
|
|
## Description
|
|
|
|
<!-- SECTION:DESCRIPTION:BEGIN -->
|
|
Integrate WebGPU-powered browser-based image generation (SD-Turbo) to reduce RunPod API costs and eliminate cold start delays. This creates a hybrid pipeline where quick drafts/iterations run locally in the browser (FREE, ~1-3 seconds), while high-quality final renders still use RunPod SDXL.
|
|
|
|
**Problem:**
|
|
- Current image generation always hits RunPod (~$0.02/image + 10-30s cold starts)
|
|
- No instant feedback loop for creative iteration
|
|
- 100% of compute costs are cloud-based
|
|
|
|
**Solution:**
|
|
- Add WebGPU capability detection
|
|
- Integrate SD-Turbo for instant browser-based previews
|
|
- Smart routing: drafts → browser, final renders → RunPod
|
|
- Potential 70% reduction in RunPod image generation costs
|
|
|
|
**Cost Impact (projected):**
|
|
- 1,000 images/mo: $20 → $6 (save $14/mo)
|
|
- 5,000 images/mo: $100 → $30 (save $70/mo)
|
|
- 10,000 images/mo: $200 → $60 (save $140/mo)
|
|
|
|
**Browser Support:**
|
|
- Chrome/Edge: Full WebGPU (v113+)
|
|
- Firefox: Windows (July 2025)
|
|
- Safari: v26 beta
|
|
- Fallback: WASM backend for unsupported browsers
|
|
<!-- SECTION:DESCRIPTION:END -->
|
|
|
|
## Acceptance Criteria
|
|
<!-- AC:BEGIN -->
|
|
- [ ] #1 WebGPU capability detection added to clientConfig.ts
|
|
- [ ] #2 SD-Turbo model loads and runs in browser via WebGPU
|
|
- [ ] #3 ImageGenShapeUtil has Quick Preview vs High Quality toggle
|
|
- [ ] #4 Smart routing in aiOrchestrator routes drafts to browser
|
|
- [ ] #5 Fallback to WASM for browsers without WebGPU
|
|
- [ ] #6 User can generate preview images with zero cold start
|
|
- [ ] #7 RunPod only called for High Quality final renders
|
|
- [ ] #8 Model download progress indicator shown to user
|
|
- [ ] #9 Works offline after initial model download
|
|
<!-- AC:END -->
|
|
|
|
## Implementation Plan
|
|
|
|
<!-- SECTION:PLAN:BEGIN -->
|
|
## Phase 1: Foundation (Quick Wins)
|
|
|
|
### 1.1 WebGPU Capability Detection
|
|
**File:** `src/lib/clientConfig.ts`
|
|
|
|
```typescript
|
|
export async function detectWebGPUCapabilities(): Promise<{
|
|
hasWebGPU: boolean
|
|
hasF16: boolean
|
|
adapterInfo?: GPUAdapterInfo
|
|
estimatedVRAM?: number
|
|
}> {
|
|
if (!navigator.gpu) {
|
|
return { hasWebGPU: false, hasF16: false }
|
|
}
|
|
|
|
const adapter = await navigator.gpu.requestAdapter()
|
|
if (!adapter) {
|
|
return { hasWebGPU: false, hasF16: false }
|
|
}
|
|
|
|
const hasF16 = adapter.features.has('shader-f16')
|
|
const adapterInfo = await adapter.requestAdapterInfo()
|
|
|
|
return {
|
|
hasWebGPU: true,
|
|
hasF16,
|
|
adapterInfo,
|
|
estimatedVRAM: adapterInfo.memoryHeaps?.[0]?.size
|
|
}
|
|
}
|
|
```
|
|
|
|
### 1.2 Install Dependencies
|
|
```bash
|
|
npm install @anthropic-ai/sdk onnxruntime-web
|
|
# Or for transformers.js v3:
|
|
npm install @huggingface/transformers
|
|
```
|
|
|
|
### 1.3 Vite Config Updates
|
|
**File:** `vite.config.ts`
|
|
- Ensure WASM/ONNX assets are properly bundled
|
|
- Add WebGPU shader compilation support
|
|
- Configure chunk splitting for ML models
|
|
|
|
---
|
|
|
|
## Phase 2: Browser Diffusion Integration
|
|
|
|
### 2.1 Create WebGPU Diffusion Module
|
|
**New File:** `src/lib/webgpuDiffusion.ts`
|
|
|
|
```typescript
|
|
import { pipeline } from '@huggingface/transformers'
|
|
|
|
let generator: any = null
|
|
let loadingPromise: Promise<void> | null = null
|
|
|
|
export async function initSDTurbo(
|
|
onProgress?: (progress: number, status: string) => void
|
|
): Promise<void> {
|
|
if (generator) return
|
|
if (loadingPromise) return loadingPromise
|
|
|
|
loadingPromise = (async () => {
|
|
onProgress?.(0, 'Loading SD-Turbo model...')
|
|
|
|
generator = await pipeline(
|
|
'text-to-image',
|
|
'Xenova/sdxl-turbo', // or 'stabilityai/sd-turbo'
|
|
{
|
|
device: 'webgpu',
|
|
dtype: 'fp16',
|
|
progress_callback: (p) => onProgress?.(p.progress, p.status)
|
|
}
|
|
)
|
|
|
|
onProgress?.(100, 'Ready')
|
|
})()
|
|
|
|
return loadingPromise
|
|
}
|
|
|
|
export async function generateLocalImage(
|
|
prompt: string,
|
|
options?: {
|
|
width?: number
|
|
height?: number
|
|
steps?: number
|
|
seed?: number
|
|
}
|
|
): Promise<string> {
|
|
if (!generator) {
|
|
throw new Error('SD-Turbo not initialized. Call initSDTurbo() first.')
|
|
}
|
|
|
|
const result = await generator(prompt, {
|
|
width: options?.width || 512,
|
|
height: options?.height || 512,
|
|
num_inference_steps: options?.steps || 1, // SD-Turbo = 1 step
|
|
seed: options?.seed
|
|
})
|
|
|
|
// Returns base64 data URL
|
|
return result[0].image
|
|
}
|
|
|
|
export function isSDTurboReady(): boolean {
|
|
return generator !== null
|
|
}
|
|
|
|
export async function unloadSDTurbo(): Promise<void> {
|
|
generator = null
|
|
loadingPromise = null
|
|
// Force garbage collection of GPU memory
|
|
}
|
|
```
|
|
|
|
### 2.2 Create Model Download Manager
|
|
**New File:** `src/lib/modelDownloadManager.ts`
|
|
|
|
Handle progressive model downloads with:
|
|
- IndexedDB caching for persistence
|
|
- Progress tracking UI
|
|
- Resume capability for interrupted downloads
|
|
- Storage quota management
|
|
|
|
---
|
|
|
|
## Phase 3: UI Integration
|
|
|
|
### 3.1 Update ImageGenShapeUtil
|
|
**File:** `src/shapes/ImageGenShapeUtil.tsx`
|
|
|
|
Add to shape props:
|
|
```typescript
|
|
type IImageGen = TLBaseShape<"ImageGen", {
|
|
// ... existing props
|
|
generationMode: 'auto' | 'local' | 'cloud' // NEW
|
|
localModelStatus: 'not-loaded' | 'loading' | 'ready' | 'error' // NEW
|
|
localModelProgress: number // NEW (0-100)
|
|
}>
|
|
```
|
|
|
|
Add UI toggle:
|
|
```tsx
|
|
<div className="generation-mode-toggle">
|
|
<button
|
|
onClick={() => setMode('local')}
|
|
disabled={!hasWebGPU}
|
|
title={!hasWebGPU ? 'WebGPU not supported' : 'Fast preview (~1-3s)'}
|
|
>
|
|
⚡ Quick Preview
|
|
</button>
|
|
<button
|
|
onClick={() => setMode('cloud')}
|
|
title="High quality SDXL (~10-30s)"
|
|
>
|
|
✨ High Quality
|
|
</button>
|
|
</div>
|
|
```
|
|
|
|
### 3.2 Smart Generation Logic
|
|
```typescript
|
|
const generateImage = async (prompt: string) => {
|
|
const mode = shape.props.generationMode
|
|
const capabilities = await detectWebGPUCapabilities()
|
|
|
|
// Auto mode: local for iterations, cloud for final
|
|
if (mode === 'auto' || mode === 'local') {
|
|
if (capabilities.hasWebGPU && isSDTurboReady()) {
|
|
// Generate locally - instant!
|
|
const imageUrl = await generateLocalImage(prompt)
|
|
updateShape({ imageUrl, source: 'local' })
|
|
return
|
|
}
|
|
}
|
|
|
|
// Fall back to RunPod
|
|
await generateWithRunPod(prompt)
|
|
}
|
|
```
|
|
|
|
---
|
|
|
|
## Phase 4: AI Orchestrator Integration
|
|
|
|
### 4.1 Update aiOrchestrator.ts
|
|
**File:** `src/lib/aiOrchestrator.ts`
|
|
|
|
Add browser as compute target:
|
|
```typescript
|
|
type ComputeTarget = 'browser' | 'netcup' | 'runpod'
|
|
|
|
interface ImageGenerationOptions {
|
|
prompt: string
|
|
priority: 'draft' | 'final'
|
|
preferLocal?: boolean
|
|
}
|
|
|
|
async function generateImage(options: ImageGenerationOptions) {
|
|
const { hasWebGPU } = await detectWebGPUCapabilities()
|
|
|
|
// Routing logic
|
|
if (options.priority === 'draft' && hasWebGPU && isSDTurboReady()) {
|
|
return { target: 'browser', cost: 0 }
|
|
}
|
|
|
|
if (options.priority === 'final') {
|
|
return { target: 'runpod', cost: 0.02 }
|
|
}
|
|
|
|
// Fallback chain
|
|
return { target: 'runpod', cost: 0.02 }
|
|
}
|
|
```
|
|
|
|
---
|
|
|
|
## Phase 5: Advanced Features (Future)
|
|
|
|
### 5.1 Real-time img2img Refinement
|
|
- Start with browser SD-Turbo draft
|
|
- User adjusts/annotates
|
|
- Send to RunPod SDXL for final with img2img
|
|
|
|
### 5.2 Browser-based Upscaling
|
|
- Add Real-ESRGAN-lite via ONNX Runtime
|
|
- 2x/4x upscale locally before cloud render
|
|
|
|
### 5.3 Background Removal
|
|
- U2Net in browser via transformers.js
|
|
- Zero-cost background removal
|
|
|
|
### 5.4 Style Transfer
|
|
- Fast neural style transfer via WebGPU shaders
|
|
- Real-time preview on canvas
|
|
|
|
---
|
|
|
|
## Technical Considerations
|
|
|
|
### Model Sizes
|
|
| Model | Size | Load Time | Generation |
|
|
|-------|------|-----------|------------|
|
|
| SD-Turbo | ~2GB | 30-60s (first) | 1-3s |
|
|
| SD-Turbo (quantized) | ~1GB | 15-30s | 2-4s |
|
|
|
|
### Memory Management
|
|
- Unload model when tab backgrounded
|
|
- Clear GPU memory on low-memory warnings
|
|
- IndexedDB for model caching (survives refresh)
|
|
|
|
### Error Handling
|
|
- Graceful degradation to WASM if WebGPU fails
|
|
- Clear error messages for unsupported browsers
|
|
- Automatic fallback to RunPod on local failure
|
|
|
|
---
|
|
|
|
## Files to Create/Modify
|
|
|
|
**New Files:**
|
|
- `src/lib/webgpuDiffusion.ts` - SD-Turbo wrapper
|
|
- `src/lib/modelDownloadManager.ts` - Model caching
|
|
- `src/lib/webgpuCapabilities.ts` - Detection utilities
|
|
- `src/components/ModelDownloadProgress.tsx` - UI component
|
|
|
|
**Modified Files:**
|
|
- `src/lib/clientConfig.ts` - Add WebGPU detection
|
|
- `src/lib/aiOrchestrator.ts` - Add browser routing
|
|
- `src/shapes/ImageGenShapeUtil.tsx` - Add mode toggle
|
|
- `vite.config.ts` - ONNX/WASM config
|
|
- `package.json` - New dependencies
|
|
|
|
---
|
|
|
|
## Testing Checklist
|
|
|
|
- [ ] WebGPU detection works on Chrome, Edge, Firefox
|
|
- [ ] WASM fallback works on Safari/older browsers
|
|
- [ ] Model downloads and caches correctly
|
|
- [ ] Generation completes in <5s on modern GPU
|
|
- [ ] Memory cleaned up properly on unload
|
|
- [ ] Offline generation works after model cached
|
|
- [ ] RunPod fallback triggers correctly
|
|
- [ ] Cost tracking reflects local vs cloud usage
|
|
<!-- SECTION:PLAN:END -->
|