| | |
| | |
| |
|
| | 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( |
| | 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) |
| |
|
