# How Image Generation Works with DALL-E 3 and Flux Schnell via Replicate in Screenshot-to-Code

> Discover how screenshot-to-code uses DALL-E 3 and Flux Schnell via Replicate for AI image generation. Learn about the direct OpenAI API and asynchronous Replicate prediction pipelines. Explore the abi/screenshot-to-code repo.

- Repository: [Abi Raja/screenshot-to-code](https://github.com/abi/screenshot-to-code)
- Tags: how-to-guide
- Published: 2026-03-02

---

**The screenshot-to-code backend handles AI image generation through two distinct pipelines—direct OpenAI API calls for DALL-E 3 and asynchronous Replicate predictions for Flux Schnell—selected dynamically via the `model` parameter in the orchestration layer.**

The `abi/screenshot-to-code` repository implements a flexible image generation system that enables AI agents to create visual assets during frontend code generation. This architecture supports both OpenAI's DALL-E 3 and Black Forest Labs' Flux models hosted on Replicate, managing API authentication, rate limiting, and provider-specific response handling. Understanding how image generation with DALL-E 3 and Flux Schnell via Replicate operates reveals the implementation details behind batch processing, asynchronous polling, and provider abstraction.

## Architecture Overview

Image generation is triggered when the `ParameterExtractionStage` detects the **`isImageGenerationEnabled`** flag (defaulting to `True`) in incoming WebSocket requests. This boolean is stored in `ExtractedParams.should_generate_images` and passed through to the agentic generation pipeline:

```python

# backend/routes/generate_code.py

should_generate_images = bool(params.get("isImageGenerationEnabled", True))

```

The core routing logic resides in **[`backend/image_generation/generation.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/image_generation/generation.py)**, where the `process_tasks` function branches based on the `model` argument:

- **`"dalle3"`** → Direct OpenAI API integration
- **`"flux"`** → Replicate API wrapper with batching and polling

## DALL-E 3 Direct Integration

When `model` equals `"dalle3"`, the system creates concurrent coroutines for each prompt using `asyncio.gather`:

```python

# backend/image_generation/generation.py

if model == "dalle3":
    tasks = [generate_image_dalle(prompt, api_key, base_url) for prompt in prompts]
    results = await asyncio.gather(*tasks, return_exceptions=True)

```

### Client Initialization and API Invocation

The `generate_image_dalle` function constructs an `AsyncOpenAI` client and invokes the **/v1/images/generations** endpoint with fixed parameters optimized for UI assets:

```python

# backend/image_generation/generation.py

client = AsyncOpenAI(api_key=api_key, base_url=base_url)
res = await client.images.generate(
    model="dall-e-3",
    quality="standard",
    style="natural",
    n=1,
    size="1024x1024",
    prompt=prompt,
)
await client.close()
return res.data[0].url if res.data else None

```

This implementation returns the direct CDN URL of the generated PNG/JPEG image, which the agent then embeds in the generated code.

## Flux Schnell via Replicate

For Flux model requests, the system implements batching to respect Replicate's rate limits. The **`REPLICATE_BATCH_SIZE`** constant (set to **20**) partitions prompts into manageable groups:

```python

# backend/image_generation/generation.py

for i in range(0, len(prompts), REPLICATE_BATCH_SIZE):
    batch = prompts[i : i + REPLICATE_BATCH_SIZE]
    tasks = [generate_image_replicate(p, api_key) for p in batch]
    results.extend(await asyncio.gather(*tasks, return_exceptions=True))

```

### Replicate API Wrapper

The `generate_image_replicate` function calls `call_replicate` in **[`backend/image_generation/replicate.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/image_generation/replicate.py)**, passing parameters compatible with the Flux model:

```python

# backend/image_generation/generation.py

return await call_replicate(
    {"prompt": prompt, "aspect_ratio": "1:1", "output_format": "png"},
    api_key,
)

```

The wrapper implements a three-step asynchronous workflow:

1. **Create Prediction**: POST to `https://api.replicate.com/v1/models/black-forest-labs/flux-2-klein-4b/predictions` with Bearer token authentication
2. **Poll for Completion**: **`MAX_POLLS = 100`** iterations with **`0.1s`** intervals until status equals `"succeeded"`
3. **Extract Output**: Normalize varied response formats (strings, dicts, or lists) into a single URL string

```python

# backend/image_generation/replicate.py

async def _run_prediction(endpoint_url: str, payload: dict, api_token: str):
    headers = {
        "Authorization": f"Bearer {api_token}",
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient() as client:
        response = await client.post(endpoint_url, headers=headers, json=payload)
        prediction_id = _extract_prediction_id(response.json())
        final_response = await _poll_prediction(client, prediction_id, headers)
        return final_response.get("output")

```

## End-to-End Execution Flow

The complete image generation lifecycle follows this sequence:

1. **WebSocket Request**: `ParameterExtractionStage` sets `should_generate_images` based on the `isImageGenerationEnabled` flag
2. **Agent Decision**: The LLM agent determines if visual assets are required for the generated code
3. **Provider Selection**: `process_tasks` routes to either `generate_image_dalle` or `generate_image_replicate` based on the configured model
4. **Async Execution**: Coroutines execute concurrently (with Replicate batches limited to 20)
5. **URL Integration**: Generated image URLs return to the agent and embed into the final HTML/Tailwind output

## Implementation Examples

### Generating Images with DALL-E 3

```python
import asyncio
from backend.image_generation.generation import process_tasks

prompts = [
    "A modern dashboard UI with dark mode and glass morphism effects",
    "Minimalist login form with soft shadows and rounded corners"
]

results = asyncio.run(
    process_tasks(
        prompts,
        api_key="sk-your-openai-key",
        base_url=None,
        model="dalle3"
    )
)
print(results)  # List of image URLs

```

### Generating Images with Flux via Replicate

```python
import asyncio
from backend.image_generation.generation import process_tasks

prompts = [
    "Hero section background with abstract gradient mesh",
    "3D icon set for e-commerce checkout flow"
]

results = asyncio.run(
    process_tasks(
        prompts,
        api_key="your-replicate-token",
        base_url=None,
        model="flux"
    )
)

```

### Low-Level Replicate Integration

For direct API access without the batching wrapper:

```python
import asyncio
import httpx

REPLICATE_API_BASE_URL = "https://api.replicate.com/v1"
FLUX_MODEL_PATH = "black-forest-labs/flux-2-klein-4b"

async def generate_flux_image(prompt: str, token: str) -> str:
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }
    payload = {
        "input": {
            "prompt": prompt,
            "aspect_ratio": "1:1",
            "output_format": "png"
        }
    }
    
    async with httpx.AsyncClient() as client:
        # Create prediction

        resp = await client.post(
            f"{REPLICATE_API_BASE_URL}/models/{FLUX_MODEL_PATH}/predictions",
            json=payload,
            headers=headers
        )
        pred_id = resp.json()["id"]
        
        # Poll until completion

        for _ in range(100):
            await asyncio.sleep(0.1)
            status = await client.get(
                f"{REPLICATE_API_BASE_URL}/predictions/{pred_id}",
                headers=headers
            )
            data = status.json()
            if data["status"] == "succeeded":
                output = data["output"]
                return output[0]["url"] if isinstance(output, list) else output
                
        raise TimeoutError("Prediction exceeded maximum polling duration")

# Usage

# url = asyncio.run(generate_flux_image("Neon cyberpunk cityscape", "token"))

```

## Summary

- **Dual-provider architecture** routes requests to either OpenAI or Replicate based on the `model` parameter in [`backend/image_generation/generation.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/image_generation/generation.py)
- **DALL-E 3 integration** uses `AsyncOpenAI` client with fixed parameters (1024x1024, standard quality, natural style) targeting the `/v1/images/generations` endpoint
- **Flux/Replicate integration** implements batching (20 prompts max) and aggressive polling (100 attempts at 0.1s intervals) via [`backend/image_generation/replicate.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/image_generation/replicate.py)
- **Model path**: Flux Schnell requests target `black-forest-labs/flux-2-klein-4b` on the Replicate platform
- **Error handling**: Both paths return `None` or exceptions for failed generations, allowing the agent to proceed without images if necessary

## Frequently Asked Questions

### How does the system choose between DALL-E 3 and Flux Schnell?

The selection occurs in `process_tasks` within [`backend/image_generation/generation.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/image_generation/generation.py) through string matching on the `model` parameter. When `model == "dalle3"`, the system invokes `generate_image_dalle`; when `model == "flux"`, it triggers the Replicate wrapper. This parameter typically originates from environment configuration or user preferences passed through the WebSocket initialization.

### Why does the Replicate implementation limit batches to 20 prompts?

The **`REPLICATE_BATCH_SIZE = 20`** constant prevents hitting Replicate's rate limits and API concurrency restrictions. Unlike OpenAI's DALL-E 3 endpoint, which handles individual requests efficiently, Replicate's prediction API requires creating and polling for each generation job separately. Batching at 20 balances throughput against API stability, while `asyncio.gather` processes each batch concurrently.

### What happens if a Replicate prediction fails or times out?

If a prediction fails, returns an error status, or exceeds the **100-poll limit** (approximately 10 seconds), the `_poll_prediction` helper raises an exception or returns an error object. The `process_tasks` function uses `return_exceptions=True` in `asyncio.gather`, allowing successful generations to complete while logging failures. The calling agent receives `None` for failed prompts and continues code generation without the asset.

### Can I use a custom OpenAI-compatible base URL for image generation?

Yes. The `generate_image_dalle` function accepts a `base_url` parameter that passes directly to the `AsyncOpenAI` client constructor. This enables integration with OpenAI-compatible APIs (such as Azure OpenAI or proxy services) by setting the base URL and providing the appropriate API key, though the endpoint must support the standard `/v1/images/generations` schema.