#  model_loader.py

import torch
from huggingface_hub import HfApi
from diffusers import (
    ControlNetModel,
    DPMSolverMultistepScheduler,
    StableDiffusionControlNetPipeline,
    StableDiffusionInpaintPipeline,
)

# Load models and preprocessor when the script starts
# Controlnet Normal
controlnet_model_id = "lllyasviel/control_v11p_sd15_normalbae"
print("initializing controlnet")
device = "cuda" if torch.cuda.is_available() else "cpu"
controlnet = ControlNetModel.from_pretrained(
    controlnet_model_id,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    attn_implementation="flash_attention_2" if torch.cuda.is_available() else None,
).to(device)

# Scheduler
scheduler = DPMSolverMultistepScheduler.from_pretrained(
    "ashllay/stable-diffusion-v1-5-archive",
    solver_order=2,
    subfolder="scheduler",
    use_karras_sigmas=True,
    final_sigmas_type="sigma_min",
    algorithm_type="sde-dpmsolver++",
    prediction_type="epsilon",
    thresholding=False,
    denoise_final=True,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
)

# Stable Diffusion Pipeline URL
base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors"
print('loading controlnet pipe')
controlnet_pipe = StableDiffusionControlNetPipeline.from_single_file(
    base_model_url,
    safety_checker=None,
    controlnet=controlnet,
    scheduler=scheduler,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
).to(device)

# Add Inpainting Pipeline Loading
# You'll need to choose an appropriate inpainting model ID
inpaint_model_id = "stabilityai/stable-diffusion-2-inpainting"
print("loading inpainting pipe")
inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
    inpaint_model_id,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
).to(device)

api = HfApi()

# Load textual inversions (only for ControlNet pipe in this version)
try:
    textual_inversions = {
        "EasyNegativeV2": "EasyNegativeV2.safetensors",
        "badhandv4": "badhandv4.pt",
        "fcNeg-neg": "fcNeg-neg.pt",
        "HDA_Ahegao": "HDA_Ahegao.pt",
        "HDA_Bondage": "HDA_Bondage.pt",
        "HDA_pet_play": "HDA_pet_play.pt",
        "HDA_unconventional_maid": "HDA_unconventional maid.pt",
        "HDA_NakedHoodie": "HDA_NakedHoodie.pt",
        "HDA_NunDress": "HDA_NunDress.pt",
        "HDA_Shibari": "HDA_Shibari.pt",
    }
    for token, weight_name in textual_inversions.items():
        try:
            controlnet_pipe.load_textual_inversion(
                "broyang/hentaidigitalart_v20", weight_name=weight_name, token=token,
            )
            print(f"Loaded textual inversion: {token}")
        except Exception as e:
            print(f"Warning: Could not load textual inversion {weight_name}: {e}")

except Exception as e:
    print(f"Error during textual inversions loading process: {e}")

print("---------------Loaded pipelines---------------")