File size: 3,566 Bytes
3d1c0e1 |
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 |
# Prediction interface for Cog ⚙️
# https://cog.run/python
import os
import subprocess
import time
import torch
from transformers import (
CLIPTextModelWithProjection,
CLIPTokenizer,
)
from diffusers import VQModel
from cog import BasePredictor, Input, Path
from src.transformer import Transformer2DModel
from src.pipeline import Pipeline
from src.scheduler import Scheduler
MODEL_CACHE = "model_cache"
MODEL_URL = (
f"https://weights.replicate.delivery/default/viiika/Meissonic/{MODEL_CACHE}.tar"
)
os.environ.update(
{
"HF_DATASETS_OFFLINE": "1",
"TRANSFORMERS_OFFLINE": "1",
"HF_HOME": MODEL_CACHE,
"TORCH_HOME": MODEL_CACHE,
"HF_DATASETS_CACHE": MODEL_CACHE,
"TRANSFORMERS_CACHE": MODEL_CACHE,
"HUGGINGFACE_HUB_CACHE": MODEL_CACHE,
}
)
def download_weights(url, dest):
start = time.time()
print("downloading url: ", url)
print("downloading to: ", dest)
subprocess.check_call(["pget", "-x", url, dest], close_fds=False)
print("downloading took: ", time.time() - start)
class Predictor(BasePredictor):
def setup(self) -> None:
"""Load the model into memory to make running multiple predictions efficient"""
if not os.path.exists(MODEL_CACHE):
download_weights(MODEL_URL, MODEL_CACHE)
model_path = f"{MODEL_CACHE}/MeissonFlow/Meissonic"
model = Transformer2DModel.from_pretrained(model_path, subfolder="transformer")
vq_model = VQModel.from_pretrained(model_path, subfolder="vqvae")
text_encoder = CLIPTextModelWithProjection.from_pretrained( # more stable sampling for some cases
f"{MODEL_CACHE}/laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
)
tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer")
scheduler = Scheduler.from_pretrained(model_path, subfolder="scheduler")
self.pipe = Pipeline(
vq_model,
tokenizer=tokenizer,
text_encoder=text_encoder,
transformer=model,
scheduler=scheduler,
).to("cuda")
def predict(
self,
prompt: str = Input(
description="Input prompt",
default="a photo of an astronaut riding a horse on mars",
),
negative_prompt: str = Input(
description="Specify things to not see in the output",
default="worst quality, low quality, low res, blurry, distortion, watermark, logo, signature, text, jpeg artifacts, signature, sketch, duplicate, ugly, identifying mark",
),
num_inference_steps: int = Input(
description="Number of denoising steps", ge=1, le=100, default=64
),
guidance_scale: float = Input(
description="Scale for classifier-free guidance", ge=0, le=20, default=9
),
seed: int = Input(
description="Random seed. Leave blank to randomize the seed", default=None
),
) -> 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}")
torch.manual_seed(seed)
image = self.pipe(
prompt=prompt,
negative_prompt=negative_prompt,
height=1024,
width=1024,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
).images[0]
output_path = f"/tmp/out.png"
image.save(output_path)
return Path(output_path)
|