Update README.md

e4cfca9 verified 5 months ago

4.28 kB

license: apache-2.0

baidu/ERNIE-Image Model Cards

Model Details

Model Description

ERNIE-Image is a text-to-image generation model developed by the ERNIE team at Baidu.

In terms of image quality, ERNIE-Image is on par with current state-of-the-art models. It demonstrates significant advantages in handling complex instructions, particularly in tasks that require accurate text rendering and knowledge-intensive generation.

Key Features

Precise text rendering: Especially strong in dense or complex text scenarios
Excellent instruction following: Accurately interprets and executes complex prompts
High-quality portraits and stylized images: Strong performance in both realism and artistic styles

Model Architecture

ERNIE-Image consists of the following components:

An 8B-parameter Diffusion Transformer (DiT)
A 3B text encoder from Ministral
A VAE based on flux2.dev
A prompt enhancer fine-tuned using Ministral 3B

Deployment

Thanks to its relatively compact model size, ERNIE-Image can be deployed on consumer-grade GPUs (e.g., 24GB VRAM), making high-quality image generation more accessible and practical.

Evaluation

Benchmark

Showcase

Uses

Installation & Download

Install the latest version of diffusers:

pip install git+https://github.com/huggingface/diffusers

Download the model:

pip install -U huggingface_hub
HF_XET_HIGH_PERFORMANCE=1 hf download baidu/ERNIE-Image

Recommended Parameters

Resolution:
- 1024x1024
- 848x1264
- 1264x848
- 768x1376
- 896x1200
- 1376x768
- 1200x896
Guidance scale: 4.0
Inference steps: 50

Usage Example

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
import random
import numpy as np
import torch
from diffusers import ErnieImagePipeline

seed = 42
print(f"seed: {seed}")
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.use_deterministic_algorithms(True)
torch.backends.cudnn.benchmark = False

# 加载 pipeline
pipe = ErnieImagePipeline.from_pretrained(
    "baidu/ERNIE-Image",
    torch_dtype=torch.bfloat16,
)
pipe = pipe.to("cuda")
pipe.transformer.eval()
pipe.vae.eval()
pipe.text_encoder.eval()
pipe.pe.eval()
# 如果是消费级显卡,例如 Nvidia 3090
# pipe.enable_model_cpu_offload()

# 设置随机种子
generator = torch.Generator(device="cuda").manual_seed(seed)
# 生成图片
output = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=5.0,
    generator=generator,
    num_images_per_prompt=1,
    use_pe=True
)
    
revised_prompt = output.revised_prompts
images = output.images
image.save(f"./hf_output_0.png")
print(revised_prompt)