fountai's picture
push
fca8815
# Prediction interface for Cog ⚙️
# https://cog.run/python
from cog import BasePredictor, Input, Path
import os
import time
import torch
import subprocess
from PIL import Image
from typing import List
from image_datasets.canny_dataset import canny_processor, c_crop
from src.flux.util import load_ae, load_clip, load_t5, load_flow_model, load_controlnet, load_safetensors
OUTPUT_DIR = "controlnet_results"
MODEL_CACHE = "checkpoints"
CONTROLNET_URL = "https://huggingface.co/XLabs-AI/flux-controlnet-canny/resolve/main/controlnet.safetensors"
T5_URL = "https://weights.replicate.delivery/default/black-forest-labs/FLUX.1-dev/t5-cache.tar"
CLIP_URL = "https://weights.replicate.delivery/default/black-forest-labs/FLUX.1-dev/clip-cache.tar"
HF_TOKEN = "hf_..." # Your HuggingFace token
def download_weights(url, dest):
start = time.time()
print("downloading url: ", url)
print("downloading to: ", dest)
subprocess.check_call(["pget", "-xf", url, dest], close_fds=False)
print("downloading took: ", time.time() - start)
def get_models(name: str, device: torch.device, offload: bool, is_schnell: bool):
t5 = load_t5(device, max_length=256 if is_schnell else 512)
clip = load_clip(device)
model = load_flow_model(name, device="cpu" if offload else device)
ae = load_ae(name, device="cpu" if offload else device)
controlnet = load_controlnet(name, device).to(torch.bfloat16)
return model, ae, t5, clip, controlnet
class Predictor(BasePredictor):
def setup(self) -> None:
"""Load the model into memory to make running multiple predictions efficient"""
t1 = time.time()
os.system(f"huggingface-cli login --token {HF_TOKEN}")
name = "flux-dev"
self.offload = False
checkpoint = "controlnet.safetensors"
print("Checking ControlNet weights")
checkpoint = "controlnet.safetensors"
if not os.path.exists(checkpoint):
os.system(f"wget {CONTROLNET_URL}")
print("Checking T5 weights")
if not os.path.exists(MODEL_CACHE+"/models--google--t5-v1_1-xxl"):
download_weights(T5_URL, MODEL_CACHE)
print("Checking CLIP weights")
if not os.path.exists(MODEL_CACHE+"/models--openai--clip-vit-large-patch14"):
download_weights(CLIP_URL, MODEL_CACHE)
self.is_schnell = False
device = "cuda"
self.torch_device = torch.device(device)
model, ae, t5, clip, controlnet = get_models(
name,
device=self.torch_device,
offload=self.offload,
is_schnell=self.is_schnell,
)
self.ae = ae
self.t5 = t5
self.clip = clip
self.controlnet = controlnet
self.model = model.to(self.torch_device)
if '.safetensors' in checkpoint:
checkpoint1 = load_safetensors(checkpoint)
else:
checkpoint1 = torch.load(checkpoint, map_location='cpu')
controlnet.load_state_dict(checkpoint1, strict=False)
t2 = time.time()
print(f"Setup time: {t2 - t1}")
def preprocess_canny_image(self, image_path: str, width: int = 512, height: int = 512):
image = Image.open(image_path)
image = c_crop(image)
image = image.resize((width, height))
image = canny_processor(image)
return image
def predict(
self,
prompt: str = Input(description="Input prompt", default="a handsome viking man with white hair, cinematic, MM full HD"),
image: Path = Input(description="Input image", default=None),
num_inference_steps: int = Input(description="Number of inference steps", ge=1, le=64, default=28),
cfg: float = Input(description="CFG", ge=0, le=10, default=3.5),
seed: int = Input(description="Random seed", default=None)
) -> List[Path]:
"""Run a single prediction on the model"""
if seed is None:
seed = int.from_bytes(os.urandom(2), "big")
print(f"Using seed: {seed}")
# clean output dir
output_dir = "controlnet_results"
os.system(f"rm -rf {output_dir}")
input_image = str(image)
img = Image.open(input_image)
width, height = img.size
# Resize input image if it's too large
max_image_size = 1536
scale = min(max_image_size / width, max_image_size / height, 1)
if scale < 1:
width = int(width * scale)
height = int(height * scale)
print(f"Scaling image down to {width}x{height}")
img = img.resize((width, height), resample=Image.Resampling.LANCZOS)
input_image = "/tmp/resized_image.png"
img.save(input_image)
subprocess.check_call(
["python3", "main.py",
"--local_path", "controlnet.safetensors",
"--image", input_image,
"--use_controlnet",
"--control_type", "canny",
"--prompt", prompt,
"--width", str(width),
"--height", str(height),
"--num_steps", str(num_inference_steps),
"--guidance", str(cfg),
"--seed", str(seed)
], close_fds=False)
# Find the first file that begins with "controlnet_result_"
for file in os.listdir(output_dir):
if file.startswith("controlnet_result_"):
return [Path(os.path.join(output_dir, file))]