| import torch | |
| from diffsynth.pipelines.flux_image_new import FluxImagePipeline, ModelConfig | |
| from PIL import Image | |
| import os | |
| import json | |
| import shutil | |
| from tqdm import tqdm | |
| for i in range(1): | |
| pipe = FluxImagePipeline.from_pretrained( | |
| torch_dtype=torch.bfloat16, | |
| device="cuda:0",# 改1 | |
| model_configs=[ | |
| ModelConfig(model_id="black-forest-labs/FLUX.1-Kontext-dev", origin_file_pattern="flux1-kontext-dev.safetensors"), | |
| ModelConfig(model_id="black-forest-labs/FLUX.1-Kontext-dev", origin_file_pattern="text_encoder/model.safetensors"), | |
| ModelConfig(model_id="black-forest-labs/FLUX.1-Kontext-dev", origin_file_pattern="text_encoder_2/"), | |
| ModelConfig(model_id="black-forest-labs/FLUX.1-Kontext-dev", origin_file_pattern="ae.safetensors"), | |
| ], | |
| ) | |
| # pipe.load_lora(pipe.dit, f"models/train/FLUX.1_lora_1203_i2i_fixcaption/epoch-{i}.safetensors", alpha=1) # 改 2 | |
| # print("load lora successfully") | |
| step = 25 | |
| input_path = "dataset/amelie_sketch_cat/epoch0" # 改3 | |
| base_path = f"dataset/amelie_sketch_cat" # 改4 | |
| save_path = f"{base_path}/validate_real_init" # 改5 | |
| # save_path_GT = f"{base_path}/GT" | |
| os.makedirs(save_path, exist_ok=True) | |
| # os.makedirs(save_path_GT, exist_ok=True) | |
| # for img in tqdm(os.listdir(input_path)): | |
| # shutil.copy2(os.path.join(input_path,img),os.path.join(save_path_GT,img)) | |
| # prompt="Convert this image into a line art style: retain the original scenes and characters unchanged, present it as a black-and-white sketch effect, and make it suitable for storyboard design. Requirements: use bold and powerful lines, highlight structures and textures with concise strokes, adopt a style close to comic sketching, roughly outline the scenes and character movements with simple lines, prohibit the depiction of details, and represent the characters' facial features with the simplest lines.", | |
| prompt = "convert this sketch into real film photos, each photo are consisted of 6 small pictures" | |
| # prompt = "Convert this image into a mbti style" | |
| print("finished copy gt") | |
| for fig in os.listdir(input_path): | |
| # if not fig.endswith(".png") or not fig.endswith(".jpg"): | |
| # continue | |
| if os.path.exists(os.path.join(save_path,fig)): | |
| continue | |
| print("begin infer") | |
| image = pipe( | |
| prompt = prompt, | |
| kontext_images=Image.open(os.path.join(input_path,fig)).resize((960, 576)), | |
| height=960, width=576, | |
| seed=0, | |
| num_inference_steps=step | |
| ) | |
| image.save(os.path.join(save_path,fig)) | |
| # for i in range(1): | |
| # pipe = FluxImagePipeline.from_pretrained( | |
| # torch_dtype=torch.bfloat16, | |
| # device="cuda", | |
| # model_configs=[ | |
| # ModelConfig(model_id="black-forest-labs/FLUX.1-Kontext-dev", origin_file_pattern="flux1-kontext-dev.safetensors"), | |
| # ModelConfig(model_id="black-forest-labs/FLUX.1-Kontext-dev", origin_file_pattern="text_encoder/model.safetensors"), | |
| # ModelConfig(model_id="black-forest-labs/FLUX.1-Kontext-dev", origin_file_pattern="text_encoder_2/"), | |
| # ModelConfig(model_id="black-forest-labs/FLUX.1-Kontext-dev", origin_file_pattern="ae.safetensors"), | |
| # ], | |
| # ) | |
| # pipe.load_lora(pipe.dit, f"models/train/FLUX.1_lora_spotlight_1203/epoch-{i}.safetensors", alpha=1) | |
| # step = 5 | |
| # base_path = f"/fi-lib/workspace/sjx/DiffSynth-Studio/validate_result/t2i_1203_STE_{step}" | |
| # save_path = f"{base_path}/epoch{i}" | |
| # os.makedirs(save_path, exist_ok=True) | |
| # with open("nano_comprehension_1201.txt", "r") as f: | |
| # prompts = f.readlines() | |
| # for prompt in prompts: | |
| # prompt = prompt.strip() | |
| # if prompt == "": | |
| # continue | |
| # prompt_dict = json.loads(prompt) | |
| # fig = f"{prompt_dict["Image_Name"]}.png" | |
| # del prompt_dict["Image_Name"] | |
| # prompt = json.dumps(prompt_dict, ensure_ascii=False) | |
| # image = pipe( | |
| # prompt ="draw a flower", | |
| # height=768, width=768, | |
| # seed=0, | |
| # num_inference_steps=step | |
| # ) | |
| # image.save(os.path.join(save_path,fig)) | |