from diffusers import AutoPipelineForText2Image
import torch

from src.linfusion import LinFusion
from src.tools import seed_everything
sd_repo = "runwayml/stable-diffusion-v1-5"

pipeline = AutoPipelineForText2Image.from_pretrained(
    sd_repo, torch_dtype=torch.float16, variant="fp16"
).to(torch.device("cuda"))
linfusion = LinFusion.construct_for(pipeline, pretrained_model_name_or_path="Yuanshi/LinFusion-1-5")
seed_everything(123)
image = pipeline(
	"An astronaut floating in space. Beautiful view of the stars and the universe in the background."
).images[0]
image.save("astronaut_space-1-5.png")