#モデルの設定を行う。

config_text = """
[FLUX]
device = auto
n_steps=24
seed=42

;model_mode = t2i
model_mode = i2i

strength = 0.6

;from_single_file = False
;base_model_path = black-forest-labs/FLUX.1-schnell
base_model_path = black-forest-labs/FLUX.1-dev

use_controlnet = False
controlnet_path = Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro

;control_mode0 = canny
;control_mode1 = tile
control_mode2 = depth
;control_mode3 = blur
;control_mode4 = pose
;control_mode5 = gray
;control_mode6 = lq

;lora_weight_repo0 = alfredplpl/flux.1-dev-modern-anime-lora
;lora_weight_path0 = modern-anime-lora-2.safetensors
;lora_weight_repo1 = aleksa-codes/flux-ghibsky-illustration
;lora_weight_path1 = lora.safetensors
;lora_scale0 = 1.0
;lora_scale1 = 1.0
;trigger_word0 = "modern anime style"
;trigger_word1 = "GHIBSKY style painting"
;lora_weight_repo2 = None
;lora_weight_path2 = ./inputs/QTCanimation_lora_v1-PAseer.safetensors
;lora_scale2 = 1.0
;trigger_word2 = ""

;select_solver = LCM
;select_solver = DPM
;select_solver = Eulera
select_solver = FMEuler

use_karras_sigmas = True
scheduler_algorithm_type = dpmsolver++
solver_order = 2

cfg_scale = 3.5
;cfg_scale = 1.0
width = 832
height = 1216
output_type = pil
aesthetic_score = 6
negative_aesthetic_score = 2.5

save_latent_simple = False
save_latent_overstep = False


"""

with open("configs/config.ini", "w", encoding="utf-8") as f:
  f.write(config_text)

class FLUX:
    def __init__(self,device = None, config_ini_path = './configs/config.ini'):

        FLUX_config = FLUXconfig(config_ini_path = config_ini_path)
        config_dict = FLUX_config.FLUX_config_dict


        if device is not None:
            self.device = device
        else:
            device = config_dict["device"]

            self.device = "cuda" if torch.cuda.is_available() else "cpu"
            if device != "auto":
                self.device = device

        self.last_latents = None
        self.last_step = -1
        self.last_timestep = 1000

        self.n_steps = int(config_dict["n_steps"])

        self.seed = int(config_dict["seed"])
        self.generator = torch.Generator(device=self.device).manual_seed(self.seed)

        self.use_controlnet = config_dict.get("use_controlnet", "False")
        if self.use_controlnet == "False":
            self.use_controlnet = False
        else:
            self.use_controlnet = True

        self.model_mode = config_dict.get("model_mode", "None")
        if self.model_mode == "t2i":
            self.model_mode = "t2i"
            print("t2i mode")
        elif self.model_mode == "i2i":
            self.model_mode = "i2i"
            print("i2i mode")
        else:
            raise ValueError("model_mode is only 't2i' or 'i2i'.")
        
        self.strength = config_dict.get("strength", "None")
        if self.strength == "None":
            self.strength = None
            if self.model_mode == "i2i":
                raise ValueError("strength is not set")
        else:
            self.strength = float(self.strength)

        self.controlnet_path = config_dict.get("controlnet_path", "None")
        if not self.use_controlnet and self.controlnet_path == "None":
            raise ValueError("controlnet_path is not set")

        self.control_modes_number_list = []
        control_mode_dict = {
            "canny":0,
            "tile":1,
            "depth":2,
            "blur":3,
            "pose":4,
            "gray":5,
            "lq":6
        }
        for i in range(7):
            if config_dict.get(f"control_mode{i}", "None") != "None":
                self.control_modes_number_list.append(control_mode_dict[config_dict[f"control_mode{i}"]])

        self.from_single_file = config_dict.get("from_single_file", "None")
        self.SINGLE_FILE_FLAG = True
        if self.from_single_file != "True":
            self.from_single_file = None
            self.SINGLE_FILE_FLAG = False


        self.base_model_path = config_dict["base_model_path"]


        self.lora_repo_list = []
        self.lora_path_list = []
        self.lora_scale_list = []
        self.lora_trigger_word_list = []
        self.lora_nums = 0
        self.LORA_FLAG = False
        for i in range(10):
            if config_dict.get(f"lora_weight_path{i}", "None") != "None":
                self.LORA_FLAG = True
                #Huggingfaceのrepoから取得する場合は、repoとpath（ファイル名）の両方を指定する
                if config_dict.get(f"lora_weight_repo{i}", "None") != "None":
                    self.lora_repo_list.append(config_dict[f"lora_weight_repo{i}"])
                #Civitaiなどからダウンロードして利用する場合は、ディレクトリ名含むpathのみで指定する
                else:
                    self.lora_repo_list.append(None)
                self.lora_path_list.append(config_dict[f"lora_weight_path{i}"])
                self.lora_nums += 1
                self.lora_scale_list.append(float(config_dict.get(f"lora_scale{i}", "1.0")))
                self.lora_trigger_word_list.append(config_dict.get(f"trigger_word{i}", "None"))


        self.select_solver = config_dict.get("select_solver", "FMEuler")

        self.use_karras_sigmas = config_dict.get("use_karras_sigmas", "True")
        if self.use_karras_sigmas == "True":
            self.use_karras_sigmas = True
        else:
            self.use_karras_sigmas = False

        self.scheduler_algorithm_type = config_dict.get("scheduler_algorithm_type", "dpmsolver++")
        self.solver_order = config_dict.get("solver_order", "None")
        if self.solver_order != "None":
            self.solver_order = int(self.solver_order)
        else:
            self.solver_order = None


        self.cfg_scale = float(config_dict["cfg_scale"])
        self.width = int(config_dict["width"])
        self.height = int(config_dict["height"])
        self.output_type = config_dict["output_type"]
        self.aesthetic_score = float(config_dict["aesthetic_score"])
        self.negative_aesthetic_score = float(config_dict["negative_aesthetic_score"])

        self.save_latent_simple = config_dict["save_latent_simple"]
        if self.save_latent_simple == "True":
            self.save_latent_simple = True
            print("use callback save_latent_simple")
        else:
            self.save_latent_simple = False

        self.save_latent_overstep = config_dict["save_latent_overstep"]
        if self.save_latent_overstep == "True":
            self.save_latent_overstep = True
            print("use callback save_latent_overstep")
        else:
            self.save_latent_overstep = False

        self.use_callback = False
        if self.save_latent_simple or self.save_latent_overstep:
            self.use_callback = True

        if (self.save_latent_simple and self.save_latent_overstep):
            raise ValueError("save_latent_simple and save_latent_overstep cannot be set at the same time")

        if self.model_mode == "t2i":
            if self.LORA_FLAG:
                self.base = self.preprepare_model_withLoRAandQuantize()
            else:
                self.base = self.preprepare_model()
        elif self.model_mode == "i2i":
            if self.LORA_FLAG:
                self.base = self.preprepare_img2img_model_withLoRAandQuantize()
            else:
                self.base = self.preprepare_img2img_model()
        else:
            raise ValueError("model_mode is only 't2i' or 'i2i'.")

class FLUX:
・・・

    def generate_image(self, prompt, neg_prompt = None, c_image_path = None, i_image_path = None, seed = None, controlnet_conditioning_scale = [1.0], temp_strength = None):
        def decode_tensors(pipe, step, timestep, callback_kwargs):
            if self.save_latent_simple:
                callback_kwargs = decode_tensors_simple(pipe, step, timestep, callback_kwargs)
            elif self.save_latent_overstep:
                callback_kwargs = decode_tensors_residual(pipe, step, timestep, callback_kwargs)
            else:
                raise ValueError("save_latent_simple or save_latent_overstep must be set")
            return callback_kwargs


        def decode_tensors_simple(pipe, step, timestep, callback_kwargs):
            latents = callback_kwargs["latents"]
            imege = None
            if self.save_latent_simple:
                image = latents_to_rgb_vae(latents,pipe)
            else:
                raise ValueError("save_latent_simple is not set")
            gettime = time.time()
            formatted_time_human_readable = time.strftime("%Y%m%d_%H%M%S", time.localtime(gettime))
            image.save(f"./outputs/latent_{formatted_time_human_readable}_{step}_{timestep}.png")

            return callback_kwargs

        def decode_tensors_residual(pipe, step, timestep, callback_kwargs):
            latents = callback_kwargs["latents"]
            if step > 0:
                residual = latents - self.last_latents
                goal = self.last_latents + residual * ((self.last_timestep) / (self.last_timestep - timestep))
                #print( ((self.last_timestep) / (self.last_timestep - timestep)))
            else:
                goal = latents

            if self.save_latent_overstep:
                image = latents_to_rgb_vae(goal,pipe)
            else:
                raise ValueError("save_latent_overstep is not set")

            gettime = time.time()
            formatted_time_human_readable = time.strftime("%Y%m%d_%H%M%S", time.localtime(gettime))
            image.save(f"./outputs/latent_{formatted_time_human_readable}_{step}_{timestep}.png")

            self.last_latents = latents
            self.last_step = step
            self.last_timestep = timestep

            if timestep == 0:
                self.last_latents = None
                self.last_step = -1
                self.last_timestep = 100

            return callback_kwargs

        def latents_to_rgb_vae(latents,pipe):

            latents = pipe._unpack_latents(latents, self.height, self.width, pipe.vae_scale_factor)
            latents = (latents / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor

            image = pipe.vae.decode(latents, return_dict=False)[0]
            image = pipe.image_processor.postprocess(image, output_type=self.output_type)

            return FluxPipelineOutput(images=image).images[0]


        def load_image_path(image_path):
            # もし image_path が str ならファイルパスとして画像を読み込む
            if isinstance(image_path, str):
                image = load_image(image_path)
                print("Image loaded from file path.")
            # もし image_path が PIL イメージならそのまま使用
            elif isinstance(image_path, Image.Image):
                image = image_path
                print("PIL Image object provided.")
            # もし image_path が Torch テンソルならそのまま使用
            elif isinstance(image_path, torch.Tensor):
                image = image_path.unsqueeze(0)
                image = image.permute(0, 3, 1, 2)
                image = image/255.0
                print("Torch Tensor object provided.")
            else:
                raise TypeError("Unsupported type. Provide a file path, PIL Image, or Torch Tensor.")

            return image

        def find_file_with_extension(image_path, i):
            # パターンに一致するファイルを検索
            file_pattern = f"{image_path}/{i}.*"
            matching_files = glob.glob(file_pattern)

            # マッチするファイルが存在する場合、そのファイルのパスを返す
            if matching_files:
                # 例: ./image_path/0.png のような完全なファイルパスが取得される
                return matching_files[0]
            else:
                # マッチするファイルがない場合
                raise FileNotFoundError(f"No file found matching pattern: {file_pattern}")
                return None

        if seed is not None:
            self.generator = torch.Generator(device=self.device).manual_seed(seed)
        if temp_strength is not None:
            self.strength = temp_strength

        control_image_list = []
        if self.use_controlnet and self.model_mode == "t2i":
            print("use controlnet mode: ",self.control_modes_number_list)
            for i in self.control_modes_number_list:
                if c_image_path is None:
                    raise ValueError("when use controlnet ,control_image_path must be set")
                control_image_name = load_image_path(find_file_with_extension(c_image_path, i))
                control_image = load_image_path(control_image_name)
                control_image_list.append(control_image)
        
        init_image = None
        if self.model_mode == "i2i":
            print("use i2i mode")
            if i_image_path is None:
                raise ValueError("when use i2i mode, init_image_path must be set")
            init_image = load_image_path(i_image_path)


        lora_weight_average = 0
        if self.LORA_FLAG:
            print("use LoRA")
            lora_weight_average = sum(self.lora_scale_list) / len(self.lora_scale_list)
            for word in self.lora_trigger_word_list:
                if (word is not None) and (word != "None"):
                    prompt = prompt + ", " + word

        image = None
        if self.model_mode == "i2i":
            if self.use_callback:
                if self.LORA_FLAG:
                    image = self.base(
                        prompt=prompt,
                        image=init_image, 
                        strength= self.strength,
                        guidance_scale=self.cfg_scale,
                        num_inference_steps=self.n_steps,
                        output_type=self.output_type,
                        width = self.width,
                        height = self.height,
                        generator=self.generator,
                        callback_on_step_end=decode_tensors,
                        callback_on_step_end_tensor_inputs=["latents"],
                        joint_attention_kwargs={"scale": lora_weight_average},
                        ).images[0]
                #LORAを利用しない場合
                else:
                    image = self.base(
                        prompt=prompt,
                        image=init_image, 
                        strength= self.strength,
                        guidance_scale=self.cfg_scale,
                        num_inference_steps=self.n_steps,
                        output_type=self.output_type,
                        width = self.width,
                        height = self.height,
                        generator=self.generator,
                        callback_on_step_end=decode_tensors,
                        callback_on_step_end_tensor_inputs=["latents"],
                        ).images[0]
            #latentを保存しない場合
            else:
                if self.LORA_FLAG:
                    image = self.base(
                        prompt=prompt,
                        image=init_image, 
                        strength= self.strength,
                        guidance_scale=self.cfg_scale,
                        num_inference_steps=self.n_steps,
                        output_type=self.output_type,
                        width = self.width,
                        height = self.height,
                        generator=self.generator,
                        joint_attention_kwargs={"scale": lora_weight_average},
                        ).images[0]
                # LORAを利用しない場合
                else:
                    image = self.base(
                        prompt=prompt,
                        image=init_image, 
                        strength= self.strength,
                        guidance_scale=self.cfg_scale,
                        num_inference_steps=self.n_steps,
                        output_type=self.output_type,
                        width = self.width,
                        height = self.height,
                        generator=self.generator
                        ).images[0]

        #t2iモードの場合
        elif self.model_mode == "t2i":
            if self.use_callback:
                if self.LORA_FLAG:
                    if self.use_controlnet:
                        image = self.base(
                            prompt=prompt,
                            control_image=control_image_list,
                            control_mode=self.control_modes_number_list,
                            guidance_scale=self.cfg_scale,
                            controlnet_conditioning_scale=controlnet_conditioning_scale,
                            num_inference_steps=self.n_steps,
                            output_type=self.output_type,
                            width = self.width,
                            height = self.height,
                            generator=self.generator,
                            callback_on_step_end=decode_tensors,
                            callback_on_step_end_tensor_inputs=["latents"],
                            joint_attention_kwargs={"scale": lora_weight_average},
                            ).images[0]
                    else:
                        image = self.base(
                            prompt=prompt,
                            guidance_scale=self.cfg_scale,
                            num_inference_steps=self.n_steps,
                            output_type=self.output_type,
                            width = self.width,
                            height = self.height,
                            generator=self.generator,
                            callback_on_step_end=decode_tensors,
                            callback_on_step_end_tensor_inputs=["latents"],
                            joint_attention_kwargs={"scale": lora_weight_average},
                            ).images[0]
                #LORAを利用しない場合
                else:
                    if self.use_controlnet:
                        image = self.base(
                            prompt=prompt,
                            control_image=control_image_list,
                            control_mode=self.control_modes_number_list,
                            guidance_scale=self.cfg_scale,
                            controlnet_conditioning_scale=controlnet_conditioning_scale,
                            num_inference_steps=self.n_steps,
                            output_type=self.output_type,
                            width = self.width,
                            height = self.height,
                            callback_on_step_end=decode_tensors,
                            callback_on_step_end_tensor_inputs=["latents"],
                            generator=self.generator
                            ).images[0]
                    else:
                        image = self.base(
                            prompt=prompt,
                            guidance_scale=self.cfg_scale,
                            num_inference_steps=self.n_steps,
                            output_type=self.output_type,
                            width = self.width,
                            height = self.height,
                            callback_on_step_end=decode_tensors,
                            callback_on_step_end_tensor_inputs=["latents"],
                            generator=self.generator
                            ).images[0]
            #latentを保存しない場合
            else:
                if self.LORA_FLAG:
                    if self.use_controlnet:
                        image = self.base(
                            prompt=prompt,
                            control_image=control_image_list,
                            control_mode=self.control_modes_number_list,
                            guidance_scale=self.cfg_scale,
                            controlnet_conditioning_scale=controlnet_conditioning_scale,
                            num_inference_steps=self.n_steps,
                            output_type=self.output_type,
                            width = self.width,
                            height = self.height,
                            generator=self.generator,
                            joint_attention_kwargs={"scale": lora_weight_average},
                            ).images[0]
                    else:
                        image = self.base(
                            prompt=prompt,
                            guidance_scale=self.cfg_scale,
                            num_inference_steps=self.n_steps,
                            output_type=self.output_type,
                            width = self.width,
                            height = self.height,
                            generator=self.generator,
                            joint_attention_kwargs={"scale": lora_weight_average},
                            ).images[0]

                # LORAを利用しない場合
                else:
                    if self.use_controlnet:
                        image = self.base(
                            prompt=prompt,
                            control_image=control_image_list,
                            control_mode=self.control_modes_number_list,
                            guidance_scale=self.cfg_scale,
                            controlnet_conditioning_scale=controlnet_conditioning_scale,
                            num_inference_steps=self.n_steps,
                            output_type=self.output_type,
                            width = self.width,
                            height = self.height,
                            generator=self.generator
                            ).images[0]
                    else:
                        image = self.base(
                            prompt=prompt,
                            guidance_scale=self.cfg_scale,
                            num_inference_steps=self.n_steps,
                            output_type=self.output_type,
                            width = self.width,
                            height = self.height,
                            generator=self.generator
                            ).images[0]
        else:
            raise ValueError("model_mode is only 't2i' or 'i2i'.")


        return image