Metadata-Version: 2.4
Name: imagegeneration
Version: 0.1.0
Summary: Automated single and high-throughput batch image generation client via Gemini web RPCs.
Keywords: image-generation,gemini,batch-processing,ai,text-to-image,multithreading
Author: TalhaZunair
Author-email: TalhaZunair <talhazunair37@gmail.com>
License-Expression: MIT
License-File: LICENSE
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Multimedia :: Graphics
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Dist: requests>=2.34.2
Requires-Python: >=3.13
Project-URL: Homepage, https://github.com/TalhaZunair/ImageGeneration
Project-URL: Repository, https://github.com/TalhaZunair/ImageGeneration
Project-URL: Issues, https://github.com/TalhaZunair/ImageGeneration/issues
Description-Content-Type: text/markdown

# imagegeneration

> A Python library for automated single and high-throughput batch image generation via Google Gemini's reverse-engineered web interface.

[![PyPI Version](https://img.shields.io/badge/pypi-v0.1.0-blue.svg)](https://pypi.org/project/imagegeneration/)
[![Python Version](https://img.shields.io/badge/python-%3E%3D3.13-blue.svg)](https://www.python.org/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Build Backend](https://img.shields.io/badge/build-uv__build-blueviolet.svg)](https://github.com/astral-sh/uv)

---

## Overview

`imagegeneration` is a Python package designed for programmatic text-to-image synthesis using Google Gemini. Unlike browser-automation frameworks (such as Selenium or Playwright) that consume significant CPU and memory resources, this library communicates directly with Gemini's underlying `batchexecute` RPC endpoint over HTTP via `requests`.

It provides two primary interfaces:
- **`ImageFoundry`**: A synchronous client for generating single images with configurable aspect ratios, retry policies, and automated directory and file naming management.
- **`ImageFoundryBatchImageGeneration`**: A multithreaded orchestrator that dispatches batches of image prompts concurrently across worker threads with live progress callbacks.

This package is intended for AI engineers, dataset creators, content pipeline developers, and automation workflows needing high-throughput image generation without browser overhead.

---

## Features

- **Direct Headless HTTP Execution**: Generates images directly via Gemini's internal `batchexecute` protocol without requiring headless Chrome, Selenium, or WebDriver installations.
- **Concurrent Batch Generation**: Employs Python's `ThreadPoolExecutor` to dispatch and process all prompt tasks simultaneously across parallel workers.
- **Dual Cookie Format Ingestion**: Natively loads session cookies from `cookies.json` formatted either as an exported list of cookie objects (`[{"name": "...", "value": "..."}]`) or a standard key-value dictionary (`{"cookie_name": "value"}`).
- **Dynamic Token Acquisition**: Automatically extracts the live anti-CSRF `SNlM0e` (`at`) parameter from `gemini.google.com/app` with a fallback mechanism and thread-safe caching.
- **Configurable Aspect Ratios**: Supports standard image aspect ratios (e.g., `16:9`, `1:1`, `9:16`, `4:3`) at both instance and per-task levels.
- **Built-in Fault Tolerance & Retries**: Automatically handles transient network glitches and RPC errors with configurable retry attempts and backoff intervals.
- **Automatic Output Organization**: Automatically creates output directories, maps response MIME types to extensions (`.jpg`, `.png`, `.webp`), and formats structured filenames.
- **Real-Time Progress Tracking**: Supports user-defined callback functions to monitor completed tasks during batch runs.

---

## Installation

### Requirements

- **Python**: `>=3.13`
- **Dependencies**: `requests>=2.34.2`

### Install via pip

```bash
pip install imagegeneration
```

### Install from Source

Clone the repository and install in editable mode:

```bash
git clone <repository-url>
cd ImageGeneration
pip install -e .
```

Or using [`uv`](https://github.com/astral-sh/uv):

```bash
uv pip install -e .
```

---

## Authentication & Credentials

This library requires an active Google Gemini web session to authenticate requests. It does not use an official Google Cloud API key; instead, it uses session cookies exported from an authenticated browser session.

### Setting Up `cookies.json`

1. Open your browser and log into [gemini.google.com](https://gemini.google.com).
2. Export your cookies for `gemini.google.com` using a browser extension (such as Cookie-Editor) or developer tools.
3. Save the exported JSON file as `cookies.json` in your project root directory (or specify a custom path in code).

#### Supported Formats

`imagegeneration` accepts either of the following structures:

**Format A: List of Cookie Objects (Default Export from extensions)**
```json
[
  {
    "name": "__Secure-1PSID",
    "value": "your_secure_1psid_cookie_here"
  },
  {
    "name": "__Secure-1PSIDTS",
    "value": "your_secure_1psidts_cookie_here"
  }
]
```

**Format B: Key-Value Dictionary**
```json
{
  "__Secure-1PSID": "your_secure_1psid_cookie_here",
  "__Secure-1PSIDTS": "your_secure_1psidts_cookie_here"
}
```

> [!WARNING]
> Keep your `cookies.json` private. Never commit it to public version control or expose session credentials. Add `cookies.json` to your `.gitignore` file.

---

## Quick Start

The simplest way to generate an image:

```python
from imagegeneration import ImageFoundry

# Initialize client (automatically loads cookies.json from current directory)
foundry = ImageFoundry(aspect_ratio="16:9")

# Generate an image
success = foundry.generate_image(
    ImagePrompt="A cinematic shot of an astronaut riding a horse on Mars at sunset",
    Output="output_images",
    image_index=1
)

if success:
    print("Image successfully generated and saved to output_images/")
else:
    print("Image generation failed.")
```

---

## Python API Reference

### `ImageFoundry`

The core synchronous client for issuing individual image generation requests.

```python
class ImageFoundry(
    cookies_path: Optional[str | Path] = None,
    cookies: Optional[Dict[str, str]] = None,
    at_token: Optional[str] = None,
    aspect_ratio: str = "16:9"
)
```

#### Parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `cookies_path` | `Optional[str \| Path]` | `None` | Path to a `cookies.json` file. Defaults to `cookies.json` in the project root if omitted. |
| `cookies` | `Optional[Dict[str, str]]` | `None` | Pre-parsed dictionary of cookies. If provided, overrides `cookies_path`. |
| `at_token` | `Optional[str]` | `None` | Pre-fetched `SNlM0e` (`at`) token. If omitted, fetched automatically from `gemini.google.com/app`. |
| `aspect_ratio` | `str` | `"16:9"` | Default aspect ratio for generated images (e.g., `"16:9"`, `"1:1"`, `"9:16"`, `"4:3"`). |

#### Methods

##### `generate_image`

Sends an image prompt to Gemini, parses the base64-encoded image response, and writes the image file to disk.

```python
def generate_image(
    self,
    ImagePrompt: str,
    Output: str = "temp/Images",
    scene_number: Optional[int] = None,
    image_index: int = 1,
    aspect_ratio: Optional[str] = None,
    max_retries: int = 3
) -> bool
```

- **Parameters**:
  - `ImagePrompt` (`str`): The descriptive text prompt for the image.
  - `Output` (`str`): Target directory path where the image will be saved. Default is `"temp/Images"`.
  - `scene_number` (`Optional[int]`): Scene identifier. When provided, the output file is named `scene{scene_number}_{image_index}.{ext}`. When `None`, the output file is named `gemini_image_{image_index}.{ext}`.
  - `image_index` (`int`): Index or iteration counter for the image. Default is `1`.
  - `aspect_ratio` (`Optional[str]`): Override aspect ratio for this specific image. If `None`, defaults to the client's `aspect_ratio`.
  - `max_retries` (`int`): Maximum number of retry attempts on network error or parsing failure. Default is `3`.
- **Returns**: `bool` — `True` if the image was successfully generated and saved; `False` otherwise.

##### `close`

```python
def close(self) -> None
```

- No-op compatibility method maintained for callers migrating from legacy Selenium-based implementations.

---

### `ImageFoundryBatchImageGeneration`

A concurrent batch orchestrator that executes multiple prompts in parallel.

```python
class ImageFoundryBatchImageGeneration(
    concurrency: Optional[int] = None,
    cookies_path: Optional[str | Path] = None,
    aspect_ratio: str = "16:9"
)
```

#### Parameters

| Parameter | Type | Default | Description |
|---|---|---|---|
| `concurrency` | `Optional[int]` | `None` | Stored concurrency setting (batch execution dynamically scales worker threads to match task volume). |
| `cookies_path` | `Optional[str \| Path]` | `None` | Path to `cookies.json`. Loaded once and shared across all parallel worker threads. |
| `aspect_ratio` | `str` | `"16:9"` | Default aspect ratio for tasks that do not define one. |

#### Methods

##### `generate_batch`

Dispatches all tasks simultaneously using a `ThreadPoolExecutor` and tracks execution.

```python
def generate_batch(
    self,
    tasks: List[Dict[str, Any]],
    output_folder: str,
    progress_callback: Optional[Callable[[int, int], None]] = None
) -> List[Dict[str, Any]]
```

- **Parameters**:
  - `tasks` (`List[Dict[str, Any]]`): List of task dictionaries. Each dictionary can contain:
    - `"prompt"` (`str`): Image prompt text (required).
    - `"scene_number"` (`Optional[int]`): Optional scene index for file naming.
    - `"image_index"` (`int`, optional): Index for filename numbering (defaults to `1`).
    - `"aspect_ratio"` (`Optional[str]`): Task-specific aspect ratio override.
  - `output_folder` (`str`): Destination directory where all generated images will be stored.
  - `progress_callback` (`Optional[Callable[[int, int], None]]`): Optional callback invoked after each task completes with `(completed_count, total_tasks)`.
- **Returns**: `List[Dict[str, Any]]` — Returns the original task dictionaries, each updated with a `"success": bool` field.

---

### Utility Functions

#### `load_cookies_from_file`

```python
def load_cookies_from_file(cookies_path: Optional[str | Path] = None) -> Dict[str, str]
```

Loads and parses cookies from a JSON file. Returns a dictionary mapping cookie names to cookie values. Returns an empty dictionary `{}` if the file does not exist or cannot be parsed.

#### `get_snlm0e_token`

```python
def get_snlm0e_token(
    cookies: Dict[str, str],
    headers: Dict[str, str],
    fallback: str = "AOvx0lJtswAqtn55NYJ19CC_GE-3:1790104383206"
) -> str
```

Fetches the live anti-CSRF `SNlM0e` (`at`) parameter by issuing a `GET` request to `https://gemini.google.com/app`. Returns the cached token if already resolved, or the fallback token if live resolution fails.

---

## Batch Image Generation

The batch generation module allows you to process multiple image generation tasks concurrently.

### Batch Task Structure

Each task dictionary supports the following keys:

```python
{
    "prompt": "Detailed description of the image",  # Required
    "scene_number": 1,                             # Optional (used in filename: scene{scene}_{index}.{ext})
    "image_index": 1,                              # Optional (defaults to 1)
    "aspect_ratio": "16:9"                         # Optional (overrides batch default)
}
```

### Complete Batch Example with Progress Callback

```python
from imagegeneration import ImageFoundryBatchImageGeneration

# Define generation tasks
tasks = [
    {
        "prompt": "Cyberpunk street market in Tokyo with neon reflections in rain puddles",
        "scene_number": 1,
        "image_index": 1,
        "aspect_ratio": "16:9"
    },
    {
        "prompt": "Macro photograph of a mechanical bee perched on a glass flower",
        "scene_number": 1,
        "image_index": 2,
        "aspect_ratio": "1:1"
    },
    {
        "prompt": "Minimalist oil painting of an alpine meadow under misty morning light",
        "scene_number": 2,
        "image_index": 1,
        "aspect_ratio": "4:3"
    }
]

# Progress reporting callback
def on_progress(completed: int, total: int):
    print(f"Progress: {completed}/{total} images completed ({(completed/total)*100:.1f}%)")

# Initialize batch generator
batch_generator = ImageFoundryBatchImageGeneration(
    cookies_path="cookies.json",
    aspect_ratio="16:9"
)

# Execute batch
results = batch_generator.generate_batch(
    tasks=tasks,
    output_folder="temp/Images",
    progress_callback=on_progress
)

# Review execution results
for item in results:
    status = "SUCCESS" if item.get("success") else "FAILED"
    print(f"[{status}] Scene {item.get('scene_number')}_{item.get('image_index')}: {item.get('prompt')[:40]}...")
```

---

## Configuration

| Option | Location | Type | Default | Description |
|---|---|---|---|---|
| `cookies_path` | `ImageFoundry.__init__` / `ImageFoundryBatchImageGeneration.__init__` | `Optional[str \| Path]` | `None` (`cookies.json`) | Filesystem path to the JSON file containing exported session cookies. |
| `cookies` | `ImageFoundry.__init__` | `Optional[Dict[str, str]]` | `None` | Directly supplied dictionary of cookie name-value pairs. |
| `at_token` | `ImageFoundry.__init__` | `Optional[str]` | `None` | Anti-CSRF `at` token. Extracted automatically from Gemini web if omitted. |
| `aspect_ratio` | Client / Method / Task | `str` | `"16:9"` | Target image aspect ratio (`"16:9"`, `"1:1"`, `"9:16"`, `"4:3"`, etc.). |
| `Output` | `generate_image` | `str` | `"temp/Images"` | Filesystem directory path for saving generated images. |
| `max_retries` | `generate_image` | `int` | `3` | Number of attempts to retry upon HTTP or parsing errors. |
| `scene_number` | `generate_image` / Task | `Optional[int]` | `None` | Scene identifier affecting filename output structure. |
| `image_index` | `generate_image` / Task | `int` | `1` | Index number for file variation naming. |

---

## Output & File Management

### Directory Structure

The destination directory is created automatically if it does not already exist:

```python
output_dir = Path(Output)
output_dir.mkdir(parents=True, exist_ok=True)
```

### File Naming Convention

File names are determined by the presence of `scene_number`:

- **With `scene_number`**: `scene{scene_number}_{image_index}.{ext}`  
  *Examples*: `scene1_1.jpg`, `scene1_2.jpg`, `scene2_1.jpg`
- **Without `scene_number`**: `gemini_image_{image_index}.{ext}`  
  *Examples*: `gemini_image_1.jpg`, `gemini_image_2.jpg`

### Format Detection

The file extension is dynamically resolved by inspecting the `mimeType` in the base64 inline data payload:

| MIME Type Returned | File Extension |
|---|---|
| `image/jpeg` / `image/jpg` | `.jpg` |
| `image/png` | `.png` |
| `image/webp` (or other) | `.webp` |

---

## Concurrency & Performance

- **Execution Model**: Batch operations leverage Python's `concurrent.futures.ThreadPoolExecutor`.
- **Parallel Dispatch**: The batch runner initializes `num_workers = max(1, len(tasks))`, submitting all tasks simultaneously.
- **Shared State**: Cookies and the `SNlM0e` (`at`) token are loaded and resolved once at initialization, allowing worker threads to share authentication data without redundant token requests.
- **Thread Safety**: Progress counters and result lists are protected by `threading.Lock` primitives.
- **Network Timeouts**: Single image requests use a 75-second socket timeout (`timeout=75`) to accommodate server-side image generation latency.

---

## CLI Usage

The package registers a command-line script entry point in `pyproject.toml`:

```toml
[project.scripts]
imagegeneration = "imagegeneration:main"
```

Currently, this entry point executes the package stub (`imagegeneration.main()`). To run image generation tasks, use the Python API described above.

---

## Error Handling & Troubleshooting

| Issue / Error Log | Cause | Solution |
|---|---|---|
| `Cookies file not found at cookies.json, using empty cookies.` | `cookies.json` is missing from the working directory. | Place your exported `cookies.json` in the root directory or pass `cookies_path="/path/to/cookies.json"`. |
| `received status code 401` / `403` | Session cookies have expired or are invalid. | Re-export active cookies from [gemini.google.com](https://gemini.google.com) and update `cookies.json`. |
| `q4uTj wrapper pattern not matched in response.` | Google has changed the RPC envelope structure or returned an error page. | Verify internet access, check that cookies are fresh, and verify your account has access to Gemini image generation. |
| `Empty candidates array.` | The model declined to generate an image for the prompt (e.g., safety or policy filters). | Review the prompt against Gemini safety guidelines and rephrase. |
| `Could not fetch live SNlM0e token` | Network failure or unexpected HTML returned by `gemini.google.com/app`. | The client will fall back to its internal default token. Ensure network connectivity. |

---

## Limitations

- **Reverse-Engineered Interface**: This package relies on Google Gemini's private web RPC (`/BardChatUi/data/batchexecute`, endpoint `q4uTj`). Google may modify, restrict, or deprecate this internal endpoint at any time without notice.
- **Cookie Lifetime**: Web session cookies expire periodically, requiring manual renewal of `cookies.json`.
- **Rate Limits & Anti-Bot Protections**: Excessive concurrent requests may trigger Google rate limits, temporary account flags, or CAPTCHA challenges.
- **Safety Filters**: Prompts triggering Gemini's trust and safety filters return empty candidate arrays and will not produce images.
- **No Official SLA**: Because this interacts with an unofficial web endpoint, uptime and response latency are subject to change.

---

## Responsible Use & Terms of Service

This package interacts with Google Gemini. Users of this library are responsible for complying with:

- Google's [Terms of Service](https://policies.google.com/terms)
- Google's [Generative AI Additional Terms of Service](https://policies.google.com/terms/generative-ai)
- Applicable laws and regulations regarding automated interactions, content generation, and intellectual property.

---

## Development & Testing

### Project Layout

```text
ImageGeneration/
├── .gitignore
├── .python-version
├── pyproject.toml
├── requirements.txt
├── uv.lock
├── README.md
├── cookies.json            # (User-provided session cookies, excluded from version control)
├── main.py                 # Core library implementation (ImageFoundry & Batch generator)
├── src/
│   └── imagegeneration/
│       └── __init__.py     # Package entrypoint (CLI stub: main())
└── temp/
    └── Images/             # Default output directory for generated images
```

### Local Development Setup

To set up a local development environment:

```bash
# Clone the repository
git clone <repository-url>
cd ImageGeneration

# Create virtual environment and install dependencies using uv
uv venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate
uv pip install -e .
```

### Testing

Automated unit tests are not currently bundled with this repository due to the requirement for live, authenticated session cookies. To test your setup, execute a single generation script with valid cookies as demonstrated in [Quick Start](#quick-start).

---

## License

This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

---

## Disclaimer

This is an independent open-source project. It is **not** an official Google product and is **not** affiliated with, maintained, sponsored, or endorsed by Google LLC or Alphabet Inc.
