Fix load pipe

[ccssu]

修复在已经加载图的情况下 触发 Input structure key None to c4612e has changed. 判断导致图重新编译。

WARNING [2024-07-10 02:37:28] /home/fengwen/worksplace/packages/onediff/src/onediff/infer_compiler/backends/oneflow/args_tree_util.py:59 - Input structure key None to c4612e has changed. Resetting the deployable module graph. This may slow down the process.
INFO [2024-07-10 02:37:28] /home/fengwen/worksplace/packages/onediff/src/onediff/infer_compiler/backends/oneflow/graph.py:14 - Building a graph for <class 'diffusers.models.autoencoders.vae.Decoder'> ...
(GRAPH:OneflowGraph_7:OneflowGraph) got a new input shape (oneflow.Size([1, 4, 128, 128]),), is creating a new graph.
(GRAPH:OneflowGraph_7:OneflowGraph) is creat

测试脚本：

python test_save_load_sdxl.py --compiler oneflow

# file: test_save_load_sdxl.py
"""
Torch run example: python examples/text_to_image_sdxl.py
Compile with oneflow: python examples/text_to_image_sdxl.py --compiler oneflow
Compile with nexfort: python examples/text_to_image_sdxl.py --compiler nexfort
Test dynamic shape: Add --run_multiple_resolutions 1 and --run_rare_resolutions 1
"""

import os
import json
import time
import argparse

import torch
import oneflow as flow

# from onediff.infer_compiler import oneflow_compile
from onediff.schedulers import EulerDiscreteScheduler
from onediffx import compile_pipe, load_pipe, save_pipe
from diffusers import StableDiffusionXLPipeline

parser = argparse.ArgumentParser()
parser.add_argument(
    "--base", type=str, default="stabilityai/stable-diffusion-xl-base-1.0"
)
parser.add_argument("--variant", type=str, default="fp16")
parser.add_argument(
    "--prompt",
    type=str,
    default="street style, detailed, raw photo, woman, face, shot on CineStill 800T",
)
parser.add_argument("--height", type=int, default=1024)
parser.add_argument("--width", type=int, default=1024)
parser.add_argument("--n_steps", type=int, default=30)
parser.add_argument("--saved_image", type=str, required=False, default="sdxl-out.png")
parser.add_argument("--seed", type=int, default=1)
parser.add_argument(
    "--compiler",
    type=str,
    default="oneflow",
    choices=["oneflow", "nexfort"],
)
parser.add_argument(
    "--compiler-config",
    type=str,
    default=None,
)
parser.add_argument(
    "--run_multiple_resolutions",
    type=(lambda x: str(x).lower() in ["true", "1", "yes"]),
    default=True,
)
parser.add_argument(
    "--run_rare_resolutions",
    type=(lambda x: str(x).lower() in ["true", "1", "yes"]),
    default=True,
)
args = parser.parse_args()

# Normal SDXL pipeline init.
OUTPUT_TYPE = "pil"

# SDXL base: StableDiffusionXLPipeline
scheduler = EulerDiscreteScheduler.from_pretrained(args.base, subfolder="scheduler")
base = StableDiffusionXLPipeline.from_pretrained(
    args.base,
    scheduler=scheduler,
    torch_dtype=torch.float16,
    variant=args.variant,
    use_safetensors=True,
)
base.to("cuda")



if args.compiler == "oneflow":
    base = compile_pipe(base)
    load_pipe(base)
    
elif args.compiler == "nexfort":
    if args.compiler_config is not None:
        options = json.loads(args.compiler_config)
    else:
        options = json.loads('{"mode": "max-autotune:cudagraphs", "dynamic": true}')

    base = compile_pipe(
        base, backend="nexfort", options=options, fuse_qkv_projections=True
    )

# Warmup with run
# Will do compilatioin in the first run
print("Warmup with running graphs...")
torch.manual_seed(args.seed)
image = base(
    prompt=args.prompt,
    height=args.height,
    width=args.width,
    num_inference_steps=args.n_steps,
    output_type=OUTPUT_TYPE,
).images

# Normal SDXL run
print("Normal SDXL run...")
torch.manual_seed(args.seed)
print(f"Running at resolution: {args.height}x{args.width}")
start_time = time.time()
image = base(
    prompt=args.prompt,
    height=args.height,
    width=args.width,
    num_inference_steps=args.n_steps,
    output_type=OUTPUT_TYPE,
).images
end_time = time.time()
print(f"Inference time: {end_time - start_time:.2f} seconds")
image[0].save(f"h{args.height}-w{args.width}-{args.saved_image}")


save_pipe(base)