File size: 1,610 Bytes
bb1a324 94bd648 bb1a324 94bd648 bb1a324 e1795a0 76c9ae3 bb1a324 | 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 | # TinyFlux Inference - Colab Cell
#
# Install and run:
# !pip install git+https://github.com/AbstractPhil/tinyflux.git
# # Then run this cell
#@title TinyFlux Inference
prompt = "object, a red apple" #@param {type:"string"}
negative_prompt = "glass, cup" #@param {type:"string"}
seed = 42 #@param {type:"integer"}
steps = 50 #@param {type:"slider", min:4, max:50, step:1}
guidance = 5 #@param {type:"slider", min:1.0, max:15.0, step:0.5}
import torch
from tinyflux.model.zoo import ModelZoo
from tinyflux.trainer.sampling import Sampler
from PIL import Image
from IPython.display import display
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32
# Load models
print("Loading models...")
zoo = ModelZoo(device=device, dtype=dtype)
zoo.load_vae()
zoo.load_clip()
zoo.load_t5()
model = zoo.load_tinyflux(
source="AbstractPhil/tiny-flux-deep",
load_ema=False, # this enables loading the checkpoint from the ema path, leave false for now to play with the most recent good version
ema_path="checkpoints/step_409244.safetensors",
)
print("✓ Models loaded")
# Create sampler
sampler = Sampler(
zoo=zoo,
model=model,
num_steps=steps,
guidance_scale=guidance,
shift=3.0,
device=device,
dtype=dtype,
)
# Generate
print(f"\nGenerating: {prompt}")
images = sampler.generate([prompt], seed=seed, negative_prompt=negative_prompt)
# Display
img = (images[0].permute(1, 2, 0).cpu().float().numpy() * 255).astype("uint8")
display(Image.fromarray(img))
print(f"Seed: {seed}, Steps: {steps}, CFG: {guidance}") |