File size: 9,537 Bytes
1e103b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 | import argparse
import base64
import io
import time
import torch
import uvicorn
import gc
import asyncio
import os
import sys
import os
import inspect
# Add OmniGen2-DFloat11 to path
# Script is in imagegen/, so we go up one level and into packages/OmniGen2-DFloat11
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
omnigen_path = os.path.join(project_root, "packages", "OmniGen2")
sys.path.insert(0, omnigen_path)
from typing import List, Optional
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from pydantic import BaseModel
from PIL import Image, ImageOps
# Import OmniGen2 and DFloat11 components
from omnigen2.pipelines.omnigen2.pipeline_omnigen2 import OmniGen2Pipeline
from omnigen2.models.transformers.transformer_omnigen2 import OmniGen2Transformer2DModel
from transformers import CLIPProcessor, BitsAndBytesConfig, Qwen2_5_VLForConditionalGeneration
from transformers.modeling_utils import no_init_weights
# Yay! Nikola here, ready to bring the OmniGen2 magic to our village!
# This server is like a new canvas for our artistic endeavors!
# Argument parsing
parser = argparse.ArgumentParser(description="OmniGen2 Image Edit Server")
parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to")
parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
# Default paths relative to project root as per plan
parser.add_argument("--base-model", type=str, default="../models/OmniGen2", help="Path to base OmniGen2 model")
parser.add_argument("--dtype", type=str, default='bf16', choices=['fp32', 'fp16', 'bf16'], help="Model precision")
args = parser.parse_args()
app = FastAPI()
# Global components
pipeline = None
request_lock = asyncio.Lock()
def load_model():
global pipeline
print(f"Loading OmniGen2 from {args.base_model}...")
# Determine usage dtype
weight_dtype = torch.float32
if args.dtype == 'fp16':
weight_dtype = torch.float16
elif args.dtype == 'bf16':
weight_dtype = torch.bfloat16
try:
# Load the base pipeline (tokenizer, scheduler, etc.)
# processor needs to be loaded separately sometimes depending on library version,
# but following inference.py pattern:
# Manually load MLLM in 4-bit to save VRAM, yay!
print("Loading MLLM in 4-bit mode for extra village efficiency!")
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=weight_dtype,
)
mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
args.base_model,
subfolder="mllm",
quantization_config=quantization_config,
torch_dtype=weight_dtype,
)
pipeline = OmniGen2Pipeline.from_pretrained(
args.base_model,
mllm=mllm,
processor=CLIPProcessor.from_pretrained(
args.base_model,
subfolder="processor",
use_fast=True
),
torch_dtype=weight_dtype,
trust_remote_code=True,
).to("cuda")
pipeline.enable_taylorseer = True
pipeline.transformer.set_attention_backend("flash")
print("Enabling CPU offload...")
#pipeline.enable_model_cpu_offload()
#pipeline.enable_sequential_cpu_offload()
except Exception as e:
print(f"Oh no! The OmniGen2 spirit refused to manifest: {e}")
raise e
print("OmniGen2 loaded successfully! Let's paint the village!")
def flush():
gc.collect()
torch.cuda.empty_cache()
class ImageGenerationRequest(BaseModel):
prompt: str
n: int = 1
size: str = "1024x1024"
response_format: str = "b64_json"
quality: str = "standard"
style: str = "vivid"
@app.on_event("startup")
async def startup_event():
load_model()
@app.post("/v1/images/edits")
async def edit_image(
image: UploadFile = File(...),
prompt: str = Form(...),
n: int = Form(1),
size: str = Form("1024x1024"),
response_format: str = Form("b64_json"),
guidance_scale: float = Form(2.5), # Image guidance scale
strength: float = Form(1.0) # Using strength to map to something or just ignored?
# OmniGen uses image_guidance_scale.
# We can map strength to text_guidance_scale maybe?
# Let's keep defaults for now from inference.py
):
if not pipeline:
raise HTTPException(status_code=500, detail="Model not loaded")
async with request_lock:
print(f"Received edit request: {prompt}")
# Processing the input image
try:
contents = await image.read()
init_image = Image.open(io.BytesIO(contents)).convert("RGB")
init_image = ImageOps.exif_transpose(init_image)
except Exception as e:
raise HTTPException(status_code=400, detail=f"Invalid image file: {e}")
# Parse max target dimensions from requested size
try:
target_width, target_height = map(int, size.split("x"))
except ValueError:
target_width, target_height = 1024, 1024
# Calculate new dimensions preserving aspect ratio
orig_width, orig_height = init_image.size
scale = min(target_width / orig_width, target_height / orig_height)
new_width = int(orig_width * scale)
new_height = int(orig_height * scale)
# Enforce multiples of 16 for compatibility
width = (new_width // 16) * 16
height = (new_height // 16) * 16
response_images = []
try:
# Generate edits
# OmniGen2Pipeline signature from inference.py:
# prompt, input_images, width, height, num_inference_steps, ...
# Using defaults from inference.py for now
results = pipeline(
prompt=prompt,
input_images=[init_image],
width=width,
height=height,
num_inference_steps=26, # Standard for OmniGen2
max_sequence_length=1024,
text_guidance_scale=5.0, # Default per inference.py
image_guidance_scale=guidance_scale, # Map guidance_scale from request here
cfg_range=(0.0, 1.0),
negative_prompt="(((deformed))), blurry, over saturation, bad anatomy, disfigured, poorly drawn face, mutation, mutated, (extra_limb), (ugly), (poorly drawn hands), fused fingers, messy drawing, broken legs censor, censored, censor_bar",
num_images_per_prompt=n,
output_type="pil",
)
for img in results.images:
buffered = io.BytesIO()
img.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
response_images.append({"b64_json": img_str})
except Exception as e:
print(f"Error during editing: {e}")
raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
finally:
flush()
return {
"created": int(time.time()),
"data": response_images
}
@app.post("/v1/images/generations")
async def generate_image(request: ImageGenerationRequest):
if not pipeline:
raise HTTPException(status_code=500, detail="Model not loaded")
async with request_lock:
print(f"Received generation request: {request.prompt}")
# Parse size
try:
width, height = map(int, request.size.split("x"))
except ValueError:
width, height = 1024, 1024
# Enforce multiples of 16 for compatibility
width = (width // 16) * 16
height = (height // 16) * 16
response_images = []
try:
# Generate images (input_images=None for txt2img)
results = pipeline(
prompt=request.prompt,
input_images=None,
width=width,
height=height,
num_inference_steps=26,
max_sequence_length=1024,
text_guidance_scale=5.0,
image_guidance_scale=2.0, # Default
cfg_range=(0.0, 1.0),
negative_prompt="(((deformed))), blurry, over saturation, bad anatomy, disfigured, poorly drawn face, mutation, mutated, (extra_limb), (ugly), (poorly drawn hands), fused fingers, messy drawing, broken legs censor, censored, censor_bar",
num_images_per_prompt=request.n,
output_type="pil",
)
for img in results.images:
buffered = io.BytesIO()
img.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
response_images.append({"b64_json": img_str})
except Exception as e:
print(f"Error during generation: {e}")
raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
finally:
flush()
return {
"created": int(time.time()),
"data": response_images
}
if __name__ == "__main__":
uvicorn.run(app, host=args.host, port=args.port)
|