Views
No views yet
1import os
2import sys
3import torch
4from kandinsky3 import get_T2I_pipeline, get_T2I_Flash_pipeline
5
6# Add kandinsky3 to Python path
7sys.path.append('..')
8
9# Set device and dtype maps
10device_map = torch.device('cuda:0')
11dtype_map = {
12 'unet': torch.float32,
13 'text_encoder': torch.float32,
14 'movq': torch.float32,
15}
16
17# Load the Flash text-to-image pipeline
18t2i_pipe = get_T2I_Flash_pipeline(
19 device_map=device_map,
20 dtype_map=dtype_map,
21 cache_dir="./cache/"
22)
23
24# Load fine-tuned UNet weights
25t2i_pipe.unet.load_state_dict(torch.load(
26 "unet_model_checkpoint.pt",
27 map_location=device_map
28))
29
30# Generate image from prompt
31res = t2i_pipe(
32 text="a new furniture design that has functions from sofa, bed, cargo, bicycle",
33 steps=50
34)[0]
35
36# Save the result
37res.save("generated_image.jpg")
38