84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
|
# model.py
|
|||
|
|
import torch
|
|||
|
|
from diffusers import FluxPipeline # 假设你用的是 flux from diffusers
|
|||
|
|
from PIL import Image
|
|||
|
|
import os
|
|||
|
|
from typing import List, Optional, Dict, Any
|
|||
|
|
|
|||
|
|
class Flux2KleinModel:
|
|||
|
|
def __init__(self, model_id="black-forest-labs/FLUX.1-dev", device="cuda"):
|
|||
|
|
print(f"Loading Flux model on {device} ...")
|
|||
|
|
self.pipe = FluxPipeline.from_pretrained(
|
|||
|
|
model_id,
|
|||
|
|
torch_dtype=torch.bfloat16
|
|||
|
|
)
|
|||
|
|
self.pipe.enable_model_cpu_offload() # 或 .to(device) 看显存情况
|
|||
|
|
# self.pipe.vae.enable_slicing() # 可选，省显存
|
|||
|
|
self.device = device
|
|||
|
|
|
|||
|
|
def preprocess_image(self, image_path: str) -> Image.Image:
|
|||
|
|
img = Image.open(image_path).convert("RGB")
|
|||
|
|
return img
|
|||
|
|
|
|||
|
|
def run_inference(
|
|||
|
|
self,
|
|||
|
|
image: Image.Image,
|
|||
|
|
prompt: str,
|
|||
|
|
aspect_ratio: str = "1:1",
|
|||
|
|
base_long_edge: int = 1024,
|
|||
|
|
steps: int = 28,
|
|||
|
|
guidance: float = 4.0,
|
|||
|
|
seed: int = 0,
|
|||
|
|
**kwargs
|
|||
|
|
) -> Image.Image:
|
|||
|
|
# 根据 aspect_ratio 计算 width/height（示例实现，可按需修改）
|
|||
|
|
if ":" in aspect_ratio:
|
|||
|
|
w, h = map(int, aspect_ratio.split(":"))
|
|||
|
|
ratio = w / h
|
|||
|
|
else:
|
|||
|
|
ratio = 1.0
|
|||
|
|
|
|||
|
|
if ratio >= 1:
|
|||
|
|
width = base_long_edge
|
|||
|
|
height = int(base_long_edge / ratio)
|
|||
|
|
else:
|
|||
|
|
height = base_long_edge
|
|||
|
|
width = int(base_long_edge * ratio)
|
|||
|
|
|
|||
|
|
# 取整到 8 的倍数（flux 常见要求）
|
|||
|
|
width = (width // 8) * 8
|
|||
|
|
height = (height // 8) * 8
|
|||
|
|
|
|||
|
|
generator = torch.Generator(device=self.device).manual_seed(seed)
|
|||
|
|
|
|||
|
|
result = self.pipe(
|
|||
|
|
prompt=prompt,
|
|||
|
|
image=image,
|
|||
|
|
height=height,
|
|||
|
|
width=width,
|
|||
|
|
num_inference_steps=steps,
|
|||
|
|
guidance_scale=guidance,
|
|||
|
|
generator=generator,
|
|||
|
|
**kwargs
|
|||
|
|
).images[0]
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
def batch_run(
|
|||
|
|
self,
|
|||
|
|
image_paths: List[str],
|
|||
|
|
prompts: List[str],
|
|||
|
|
out_paths: List[str],
|
|||
|
|
**common_kwargs
|
|||
|
|
):
|
|||
|
|
assert len(image_paths) == len(prompts) == len(out_paths)
|
|||
|
|
|
|||
|
|
for img_path, prompt, out_path in zip(image_paths, prompts, out_paths):
|
|||
|
|
try:
|
|||
|
|
img = self.preprocess_image(img_path)
|
|||
|
|
result = self.run_inference(img, prompt, **common_kwargs)
|
|||
|
|
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
|||
|
|
result.save(out_path)
|
|||
|
|
print(f"Saved: {out_path}")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Failed {img_path}: {e}")
|