import requests
import json
import os
import sys
from PIL import Image

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(os.path.dirname(__dir__)))

from modules.tools.image_utils import image2base64, base642image

data = {
    "prompt": "a pretty puppy", 
    "negative_prompt": "ugly, low quaility, blur",
    "seed": 2023
}
# fill in your ip address and port like 0.0.0.0:1001
response = requests.post('http://x.x.x.x:x/text2image', json=data)
result = json.loads(response.content.decode("utf-8"))
base64_str = result["images"][0]
image = base642image(base64_str)
image.save("outputs/agent_txt2img.jpg")
backend/agents/txt2img.py
MooreThreads/MobiMaliangSDK
import torch
import uvicorn
import os
import sys
import time
from fastapi import File, UploadFile, Request, FastAPI
from fastapi import Request
import json

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(os.path.dirname(__dir__)))

from modules.shared import SDModel, device
from modules.basic.txt2img import txt2img

from modules.tools.image_utils import image2base64, base642image
from modules.tools.loggers import server_logger

MODEL_DIR = "models"
BASE_MODEL_DIR = os.path.join(MODEL_DIR, "Stable-diffusion/ml-person-1.2.safetensors")
LORA_PATH = os.path.join(MODEL_DIR, "models/lora_models/blindbox_v1_mix.safetensors")

sd_model = SDModel(device=device, requires_safety_checker=False)
sd_model.load_models(BASE_MODEL_DIR, diffusers_format=False)
app = FastAPI()

@app.post("/text2image")
async def text2image(request: Request):
    try:
        params = await request.json()
        prompt = params.pop("prompt", "masterpiece, best quaility")
        seed = params.pop("seed", 0)
        images, status = txt2img(
            sd_model,
            prompt=prompt,
            seed=seed,
            **params
        )
        decoded_images = [image2base64(image) for image in images]
        rsp = {"images": decoded_images, "message":"ok"}
        
    except Exception as e:
        server_logger.error(" ------ {}".format(e))
        rsp = {"images": [], "message": f"err: {e}"}
    
    return rsp

@app.post("/lora2image")
async def lora2image(request: Request):
    try:
        params = await request.json()
        lora_path = params.pop("lora_path", LORA_PATH)
        sd_model.load_lora(lora_path, alpha=1)
        assert lora_path is not None
        
        prompt = params.pop("prompt", "masterpiece, best quaility")
        seed = params.pop("seed", 0)
        images, status  = txt2img(
            sd_model,
            prompt,
            seed=seed,
            lora_model_name_or_path_or_dict=LORA_PATH,
            lora_alpha=0.75,
            **params
        )
        decoded_images = [image2base64(image) for image in images]
        
        rsp = {"images": decoded_images, "message":"ok"}
    except Exception as e:
        server_logger.error(" ------ {}".format(e))
        rsp = {"images": [], "message": f"err: {e}"}
    return rsp

if __name__ == '__main__':
    uvicorn.run(app, port=1001, host="0.0.0.0")
backend/server.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import CannyDetectorWrapper
from modules.controlnet.canny2img import canny2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--canny_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the canny controlnet model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = CannyDetectorWrapper()
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.canny_controlnet_path, diffusers_format=True)

    images, _, canny_image = canny2img(
        sd_model,
        controlnet_model,
        processor,
        ori_image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "canny2img.png"))
    canny_image.save(os.path.join(args.output_path, "canny.png"))
# python examples/controlnets/example_shuffle2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --shuffle_controlnet_path models/controlnet/shuffle_v11 --image_path data/person.jpg --prompt "a pretty girl" --output_path outputs
examples/controlnets/example_canny2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import InpaintingConditionalImageWrapper
from modules.controlnet.inpaint2img import inpaint2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--inpainting_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the inpainting controlnet model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )
    
    parser.add_argument(
        "--mask_image_path", default=None, type=str, required=True, help="The mask image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    mask_iamge = Image.open(args.mask_image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = InpaintingConditionalImageWrapper()
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.inpainting_controlnet_path, diffusers_format=True)

    images, _, inpainting_image = inpaint2img(
        sd_model,
        controlnet_model,
        processor,
        ori_image=image,
        mask_image=mask_iamge,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "inpaint2img.png"))
    inpainting_image.save(os.path.join(args.output_path, "inpaint.png"))
    
# python examples/controlnets/example_inpaint2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --inpainting_controlnet_path models/controlnet/inpaint_v11 --image_path data/original.png --mask_image_path data/mask.png  --prompt "a pretty rabbit" --output_path outputs
examples/controlnets/example_inpaint2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.controlnet.ip2p2img import ip2p2imgdevice

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--ip2p_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the ip2p controlnet model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.ip2p_controlnet_path, diffusers_format=True)

    images, _, ip2p_image = ip2p2img(
        sd_model,
        controlnet_model,
        ori_image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "ip2p2img.png"))
    ip2p_image.save(os.path.join(args.output_path, "ip2p.png"))
    
# python examples/controlnets/example_ip2p2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --ip2p_controlnet_path models/controlnet/ip2p_v11 --image_path data/person.jpg --prompt "a pretty girl" --output_path outputs
    
examples/controlnets/example_ip2p2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import LineartDetectorWrapper
from modules.controlnet.lineart2img import lineart2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--lineart_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the lineart controlnet model.",
    )
    
    parser.add_argument(
        "--lineart_processor_model_path",
        default=None,
        type=str,
        required=True,
        help="Path to the lineart processor model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = LineartDetectorWrapper(device=sd_model.device)
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.lineart_controlnet_path, diffusers_format=True)
    processor.load_models(args.lineart_processor_model_path)

    images, _, lineart_image = lineart2img(
        sd_model,
        controlnet_model,
        processor,
        ori_image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "lineart2img.png"))
    lineart_image.save(os.path.join(args.output_path, "lineart.png"))

# python examples/controlnets/example_lineart2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --lineart_controlnet_path models/controlnet/lineart_v11 --lineart_processor_model_path models/controlnet/annotators --image_path data/person.jpg --prompt "a pretty girl" --output_path outputs

examples/controlnets/example_lineart2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import MLSDdetectorWrapper
from modules.controlnet.mlsd2img import mlsd2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--mlsd_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the mlsd controlnet model.",
    )
    
    parser.add_argument(
        "--mlsd_processor_model_path",
        default=None,
        type=str,
        required=True,
        help="Path to the mlsd processor model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = MLSDdetectorWrapper(device="cpu")
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.mlsd_controlnet_path, diffusers_format=True)
    processor.load_models(args.mlsd_processor_model_path)

    images, _, mlsd_image = mlsd2img(
        sd_model,
        controlnet_model,
        processor,
        ori_image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "mlsd2img.png"))
    mlsd_image.save(os.path.join(args.output_path, "mlsd.png"))

# python examples/controlnets/example_mlsd2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --mlsd_controlnet_path models/controlnet/mlsd_v11 --mlsd_processor_model_path models/controlnet/annotators --image_path data/person.jpg --prompt "a pretty girl" --output_path outputs

examples/controlnets/example_mlsd2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import NormalBaeDetectorWrapper
from modules.controlnet.normalbae2img import normalbae2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--normalbae_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the normalbae controlnet model.",
    )
    
    parser.add_argument(
        "--normalbae_processor_model_path",
        default=None,
        type=str,
        required=True,
        help="Path to the normalbae processor model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = NormalBaeDetectorWrapper("cpu")
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.normalbae_controlnet_path, diffusers_format=True)
    processor.load_models(args.normalbae_processor_model_path)

    images, _, normalbae_image = normalbae2img(
        sd_model,
        controlnet_model,
        processor,
        ori_image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "normalbae2img.png"))
    normalbae_image.save(os.path.join(args.output_path, "normalbae.png"))

# python examples/controlnets/example_normalbae2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --normalbae_controlnet_path models/controlnet/normalbae_v11 --normalbae_processor_model_path models/controlnet/annotators --image_path data/person.jpg --prompt "a pretty girl" --output_path outputs
examples/controlnets/example_normalbae2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import OpenposeDetectorWrapper
from modules.controlnet.pose2img import pose2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--openpose_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the openpose controlnet model.",
    )
    
    parser.add_argument(
        "--openpose_processor_model_path",
        default=None,
        type=str,
        required=True,
        help="Path to the openpose processor model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = OpenposeDetectorWrapper(sd_model.device)
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.openpose_controlnet_path, diffusers_format=True)
    processor.load_models(args.openpose_processor_model_path)

    images, _, openpose_image = pose2img(
        sd_model,
        controlnet_model,
        processor,
        ori_image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "openpose2img.png"))
    openpose_image.save(os.path.join(args.output_path, "openpose.png"))

# python examples/controlnets/example_pose2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --openpose_controlnet_path models/controlnet/openpose_v11 --openpose_processor_model_path models/controlnet/annotators --image_path data/person.jpg --prompt "a pretty girl" --output_path outputs
examples/controlnets/example_pose2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import HEDdetectorWrapper
from modules.controlnet.scribble2img import scribble2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--scribble_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the scribble controlnet model.",
    )
    
    parser.add_argument(
        "--scribble_processor_model_path",
        default=None,
        type=str,
        required=True,
        help="Path to the scribble processor model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = HEDdetectorWrapper(device=sd_model.device)
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.scribble_controlnet_path, diffusers_format=True)
    processor.load_models(args.scribble_processor_model_path)

    images, _, scribble_image = scribble2img(
        sd_model,
        controlnet_model,
        processor,
        ori_image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "scribble2img.png"))
    scribble_image.save(os.path.join(args.output_path, "scribble.png"))

# python examples/controlnets/example_scribble2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --scribble_controlnet_path models/controlnet/scribble_v11 --scribble_processor_model_path models/controlnet/annotators --image_path data/person.jpg --prompt "a pretty girl" --output_path outputs
examples/controlnets/example_scribble2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import ContentShuffleDetectorWrapper
from modules.controlnet.shuffle2img import shuffle2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--shuffle_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the shuffle controlnet model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = ContentShuffleDetectorWrapper()
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.shuffle_controlnet_path, diffusers_format=True)

    images, _, shuffle_image = shuffle2img(
        sd_model,
        controlnet_model,
        processor,
        ori_image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "shuffle2img.png"))
    shuffle_image.save(os.path.join(args.output_path, "shuffle.png"))
    
    # python examples/controlnets/example_shuffle2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --shuffle_controlnet_path models/controlnet/shuffle_v11 --image_path data/person.jpg --prompt "a pretty girl" --output_path outputs
examples/controlnets/example_shuffle2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import PidiNetDetectorWrapper
from modules.controlnet.softedge2img import softedge2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--softedge_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the softedge controlnet model.",
    )
    
    parser.add_argument(
        "--softedge_processor_model_path",
        default=None,
        type=str,
        required=True,
        help="Path to the softedge processor model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = PidiNetDetectorWrapper(device="cpu")
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.softedge_controlnet_path, diffusers_format=True)
    processor.load_models(args.softedge_processor_model_path)

    images, _, softedge_image = softedge2img(
        sd_model,
        controlnet_model,
        processor,
        ori_image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "softedge2img.png"))
    softedge_image.save(os.path.join(args.output_path, "softedge.png"))

# python examples/controlnets/example_softedge2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --softedge_controlnet_path models/controlnet/softedge_v11 --softedge_processor_model_path models/controlnet/annotators --image_path data/person.jpg --prompt "a pretty girl" --output_path outputs
examples/controlnets/example_softedge2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import ResizeForConditionalImageWrapper
from modules.controlnet.tile2img import tile2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--tile_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the tile controlnet model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = ResizeForConditionalImageWrapper()
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.tile_controlnet_path, diffusers_format=True)

    images, _, tile_image = tile2img(
        sd_model,
        controlnet_model,
        processor,
        ori_image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "tile2img.png"))
    tile_image.save(os.path.join(args.output_path, "tile.png"))
    
    # python examples/controlnets/example_tile2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --tile_controlnet_path models/controlnet/tile_v11 --image_path data/person.jpg --prompt "a pretty girl" --output_path outputs
examples/controlnets/example_tile2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet, device
from modules.wrappers import ZoomConditionalImageWrapper
from modules.controlnet.tilehd2img import tilehd2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--tile_controlnet_path",
        default=None,
        type=str,
        required=True,
        help="Path to the tile controlnet model.",
    )

    parser.add_argument(
        "--image_path", default=None, type=str, required=True, help="The image path"
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    controlnet_model = Controlnet(device=sd_model.device)
    processor = ZoomConditionalImageWrapper()
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    controlnet_model.load_models(args.tile_controlnet_path, diffusers_format=True)
    
    images, _, zoom_image = tilehd2img(
        sd_model, 
        controlnet_model, 
        processor, 
        prompt=args.prompt, 
        negative_prompt=args.negative_prompt,
        ori_image=image, 
        up_sampling_ratio=2,
        window_size=80,
        strength=0.35,
    )
    images[0].save("outputs/tilehd2img.png")
    zoom_image.save("outputs/zoom.png")
    
    # python examples/controlnets/example_tilehd2img.py --checkpoint_path models/Stable-diffusion/v1-5-pruned-emaonly.ckpt --tile_controlnet_path models/controlnet/tile_v11 --image_path data/person.jpg --prompt "(best quality, ultra detailed, master piece, sharp:1.3)" --output_path outputs

examples/controlnets/example_tilehd2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.shared import SDModel, device
from modules.basic.img2img import img2img
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--image_path",
        default=None,
        type=str,
        required=True,
        help="The initial image path",
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    image = Image.open(args.image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)

    images, _ = img2img(
        sd_model,
        image=image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "img2img.png"))

examples/example_img2img.py
MooreThreads/MobiMaliangSDK
import torch
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.shared import SDModel, device
from modules.basic.txt2img import txt2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--lora_path",
        default=None,
        type=str,
        required=True,
        help="Path to the Lora model.",
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--lora_alpha", default=0.75, required=False, type=float, help="Lora weight"
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    sd_model = SDModel(device=device)
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)

    image, status = txt2img(
        sd_model,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
        lora_model_name_or_path_or_dict=args.lora_path,
        lora_alpha=args.lora_alpha,
    )

    os.makedirs(args.output_path, exist_ok=True)
    image[0].save(os.path.join(args.output_path, "lora2img.png"))

examples/example_lora2img.py
MooreThreads/MobiMaliangSDK
import torch
from PIL import Image
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.shared import SDModel, device
from modules.reference.reference import reference2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--ref_image_path",
        default=None,
        type=str,
        required=True,
        help="The reference image path",
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    ref_image = Image.open(args.ref_image_path).convert("RGB").resize((512, 512))
    sd_model = SDModel(device=device)
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)
    images, _ = reference2img(
        sd_model,
        ref_image=ref_image,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "reference2img.png"))

examples/example_reference2img.py
MooreThreads/MobiMaliangSDK
import torch
import os
import sys
from PIL import Image
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.wrappers import RealESRGANWrapper
from modules.shared import device

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to super resolution model.",
    )

    parser.add_argument(
        "--image_path",
        default=None,
        type=str,
        required=True,
        help="The input image path",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    args = parser.parse_args()

    realesrgan = RealESRGANWrapper(device=device)
    realesrgan.load_models(model_path=args.checkpoint_path)
    image = Image.open(args.image_path).convert("RGB")
    result = realesrgan(image)

    os.makedirs(args.output_path, exist_ok=True)
    result.save(os.path.join(args.output_path, "sr_img.png"))

examples/example_super_resolution.py
MooreThreads/MobiMaliangSDK

import os
import sys
import torch
from PIL import Image
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.wrappers import DeepDanbooruWrapper
from modules.shared import device

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to super resolution model.",
    )

    parser.add_argument(
        "--image_path",
        default=None,
        type=str,
        required=True,
        help="The input image path",
    )

    args = parser.parse_args()

    tagger = DeepDanbooruWrapper(device=device)
    tagger.load_models(model_path=args.checkpoint_path)
    image = Image.open(args.image_path).convert("RGB")
    result = tagger(image)
    print(result)

examples/example_tagger.py
MooreThreads/MobiMaliangSDK
import torch
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.shared import SDModel, device
from modules.basic.txt2img import txt2img
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--text_inversion_path",
        default=None,
        type=str,
        required=True,
        help="Path to the text inventrual model.",
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )
    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    sd_model = SDModel(device=device)
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)

    image, status = txt2img(
        sd_model,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
        text_inversion_model_name_or_path=args.text_inversion_path,
    )

    os.makedirs(args.output_path, exist_ok=True)
    image[0].save(os.path.join(args.output_path, "textinventual2img.png"))

examples/example_textinversion2img.py
MooreThreads/MobiMaliangSDK

import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.wrappers import Translator

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the translator model.",
    )

    parser.add_argument(
        "--chinese_text",
        default=None,
        type=str,
        required=True,
        help="Chinese sentence(s) or phrase(s) or word(s)",
    )

    args = parser.parse_args()

    tagger = Translator(device="cpu")
    tagger.load_models(model_path=args.checkpoint_path)
    english = tagger(args.chinese_text)
    print(english)

examples/example_translator.py
MooreThreads/MobiMaliangSDK
import torch
import os
import sys
import argparse

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.shared import SDModel, device
from modules.basic.txt2img import txt2img

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the stable diffusion model.",
    )

    parser.add_argument(
        "--prompt",
        default=None,
        type=str,
        required=True,
        help="Describe the image content with text",
    )

    parser.add_argument(
        "--output_path",
        default="./",
        type=str,
        required=True,
        help="Result folder path",
    )

    parser.add_argument(
        "--negative_prompt",
        default=None,
        type=str,
        required=False,
        help="Content which you don't want to show up with text",
    )

    parser.add_argument(
        "--diffusers_format",
        action="store_true",
        required=False,
        help="Model files in diffusers format",
    )

    parser.add_argument(
        "--seed", default=2023, required=False, type=int, help="Random seed"
    )

    args = parser.parse_args()

    sd_model = SDModel(device=device, requires_safety_checker=False)
    sd_model.load_models(args.checkpoint_path, diffusers_format=args.diffusers_format)

    images, status = txt2img(
        sd_model,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
    )

    os.makedirs(args.output_path, exist_ok=True)
    images[0].save(os.path.join(args.output_path, "txt2img.png"))

examples/example_txt2img.py
MooreThreads/MobiMaliangSDK
NEED_SAFETY = False
SAMPLER = "DDIM"
MAX_SIZE = 768
DEBUG = True
FREE_HW = True

frontend/config.py
MooreThreads/MobiMaliangSDK
# -*- coding: UTF-8 -*-

# 摩笔马良WebUI前端代码
# 包括模型处理, 页面样式, 侧边栏, 创作tab, 画廊tab, 工具tab六部分内容

import numpy as np
import streamlit as st
from PIL import Image, ImageFilter, PngImagePlugin
from glob import glob
from tqdm import tqdm
import pickle
import os
import cv2
import random
import time
import sys
from io import BytesIO
from config import *
import warnings
warnings.filterwarnings("ignore")

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))

import torch
import torch.nn.functional as F
from safetensors.torch import load_file
from diffusers import logging
logging.set_verbosity_error()

from modules.shared import SDModel, Controlnet, MTGPU_DETECTION, device
from modules.basic.txt2img import txt2img
from modules.basic.img2img import img2img
from modules.controlnet.canny2img import canny2img
from modules.controlnet.pose2img import pose2img
from modules.controlnet.tilehd2img import tilehd2img
from modules.controlnet.mlsd2img import mlsd2img
from modules.tools.image_utils import read_img_as_base64, read_png_info
from modules.tools.prompt_parser import get_lora_from_prompt, contain_zh
from modules.tools.prompt_engineering import get_enhancer, get_artists, get_neg_prompt
from modules.wrappers import *

from transformers import CLIPModel

##### maliang #####

args, btns = {}, {}

args["save_dir"] = os.path.join("frontend", "output")
os.makedirs(args["save_dir"], exist_ok=True)

args["sd_ckpts"] = sorted(glob(os.path.join("models", "Stable-diffusion", "*")))
args["sd_ckpts"] = [path.split(os.sep)[-1] for path in args["sd_ckpts"]]

args["text_embeddings"] = sorted(glob(os.path.join("models", "embeddings", "*.pt")))

args["lora_names"] = sorted(glob(os.path.join("models", "lora", "*.safetensors")))
args["lora_names"] = [
    "<lora:" + path.split(os.sep)[-1][:-len(".safetensors")] + ":0.8>"
    for path in args["lora_names"]
]
args["loras"] = {}

random_prompt_path = np.random.choice(
    glob(os.path.join("assets", "random_prompts", "*"))
)
with open(random_prompt_path, "r") as fr:
    lines = fr.readlines()
    args["random_prompts"] = [line.rstrip("\n") for line in lines]

args["sampler"] = SAMPLER

args["loaded_loras"] = os.path.join("models", "loaded_loras")
os.makedirs(args["loaded_loras"], exist_ok=True)
os.makedirs(os.path.join(args["loaded_loras"], "fixed"), exist_ok=True)
os.makedirs(os.path.join(args["loaded_loras"], "dynamic"), exist_ok=True)
loaded_ckpt_path = os.path.join(args["loaded_loras"], "loaded_ckpt.txt")


@st.cache_resource(max_entries=1)
def load_ckpt():
    ml = SDModel(device=device, requires_safety_checker=NEED_SAFETY)

    path = os.path.join("models", "Stable-diffusion", st.session_state.ckpt_name)
    ml.load_models(path, diffusers_format=os.path.isdir(path))

    _, _ = txt2img(
        sd_model=ml,
        prompt="",
        height=256,
        width=256,
        num_inference_steps=1,
        guidance_scale=7.5,
        negative_prompt=None,
        num_images_per_prompt=1,
        seed=1234,
        callback=None,
        text_inversion_model_name_or_path=args["text_embeddings"],
        fast_mode=False,
    )

    with open(loaded_ckpt_path, "w") as fw:
        fw.write(st.session_state.ckpt_name)

    return ml

def switch_ckpt():
    ml.offload_models()

    path = os.path.join("models", "Stable-diffusion", st.session_state.ckpt_name)
    ml.load_models(path, diffusers_format=os.path.isdir(path))

    _, _ = txt2img(
        sd_model=ml,
        prompt="",
        height=256,
        width=256,
        num_inference_steps=1,
        guidance_scale=7.5,
        negative_prompt=None,
        num_images_per_prompt=1,
        seed=1234,
        callback=None,
        text_inversion_model_name_or_path=args["text_embeddings"],
        fast_mode=False,
    )

    with open(loaded_ckpt_path, "w") as fw:
        fw.write(st.session_state.ckpt_name)

@st.cache_resource()
def load_controlnet(control_type):
    net = Controlnet(device=device)

    if control_type == "canny":
        processor = CannyDetectorWrapper()
        net.load_models(
            os.path.join("models", "controlnet", "canny_v11"), diffusers_format=True
        )

    elif control_type == "pose":
        processor = OpenposeDetectorWrapper(device=device)
        processor.load_models(os.path.join("models", "controlnet", "annotators"))
        net.load_models(
            os.path.join("models", "controlnet", "openpose_v11"), diffusers_format=True
        )

    elif control_type == "tile":
        processor = ZoomConditionalImageWrapper()
        net.load_models(
            os.path.join("models", "controlnet", "tile_v11"), diffusers_format=True
        )

    elif control_type == 'mlsd':
        processor = MLSDdetectorWrapper(device="cpu")
        processor.load_models(os.path.join("models", "controlnet", "annotators"))
        net.load_models(
            os.path.join("models", "controlnet", "mlsd_v11"), diffusers_format=True
        )

    return processor, net


mlcn_canny, mlcn_pose, mlcn_tile, mlcn_mlsd = None, None, None, None


@st.cache_resource()
def load_safetensors(safetensors_name):
    sf = load_file(safetensors_name)
    return sf


@st.cache_resource()
def load_translator():
    ml_translator = Translator(device="cpu")
    ml_translator.load_models(os.path.join("models", "tools", "zh2en"))
    return ml_translator


ml_translator = None


@st.cache_resource()
def load_sr():
    ml_sr = RealESRGANWrapper(device=device)
    ml_sr.load_models(os.path.join("models", "tools", "RealESRGAN_x4plus.pth"))
    return ml_sr


ml_sr = None


@st.cache_resource()
def load_tagger():
    ml_tagger = DeepDanbooruWrapper(device=device)
    ml_tagger.load_models(os.path.join("models", "tools", "model-resnet_custom_v3.pt"))
    return ml_tagger


ml_tagger = None


@st.cache_resource()
def load_enhancer():
    with open(os.path.join("assets", "enhancer.pkl"), "rb") as fr:
        data_enhancer = pickle.load(fr)

    with open(os.path.join("assets", "artist.txt"), "r") as fr:
        lines = fr.readlines()
        artists_list = [line.rstrip("\n").split("^")[0] for line in lines]

    return [data_enhancer, artists_list]


ml_enhancer = load_enhancer()


@st.cache_resource
def load_dynamic_loras():
    dynamic_loras_dir = os.path.join("models", "tools", "dynamic_loras")

    if os.path.exists(dynamic_loras_dir):
        ml_clip = CLIPModel.from_pretrained(os.path.join("models", "tools", "clip-vit-large-patch14"))
        ml_cc = np.load(os.path.join(dynamic_loras_dir, "centroids.npy"))
        ml_loras_paths = sorted(glob(os.path.join(dynamic_loras_dir, "lora", "*.safetensors")))
    else:
        ml_clip, ml_cc, ml_loras_paths = None, None, None

    return ml_clip, ml_cc, ml_loras_paths

ml_clip, ml_cc, ml_loras_paths = load_dynamic_loras()


##### style #####

st.markdown(
    """
    <style type='text/css'>
    .css-pxxe24 {
        visibility: hidden;
    }
    div[data-testid='stTickBar'] {
        display: none;
    }
    div.block-container {
        padding-bottom: 50px;
    }

    section[data-testid='stSidebar'] {
        width: 25% !important;
        min-width: 356px !important;
    }
    section[data-testid='stSidebar'] div.block-container {
        padding-bottom: 0px;
    }
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(1) {
        display: none;
    }
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) {
        padding: 1.8rem;
    }
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(1) p {
        margin-top: -10px;
    }
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(1) div[role='progressbar'] {
        margin-top: 10px;
        margin-bottom: 20px;
    }
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(2),
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(3) {

    }
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(2) > div[data-testid='stVerticalBlock'] div[data-testid='stImage'],
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(3) > div[data-testid='stVerticalBlock'] div[data-testid='stImage'] {
        width: 100%;
    }
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(2) > div[data-testid='stVerticalBlock'] button[title='View fullscreen'],
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(3) > div[data-testid='stVerticalBlock'] button[title='View fullscreen'] {
        display: none;
    }
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(2) > div[data-testid='stVerticalBlock'] button[kind='secondary'],
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(3) > div[data-testid='stVerticalBlock'] button[kind='secondary'],
    div.st-bc > div:nth-child(5) div.stDownloadButton button[kind='secondary'] {
        position: absolute;
        top: -60px;
        right: 7px;
        background-color: rgba(13, 17, 23, 0.4);
        width: 80px;
        border-radius: 4px;
    }
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(2) > div[data-testid='stVerticalBlock'] button[kind='secondary']:hover,
    section[data-testid='stSidebar'] > div:first-child > div:nth-child(2) > div:first-child > div:first-child > div:first-child > div:nth-child(3) > div[data-testid='stVerticalBlock'] button[kind='secondary']:hover,
    div.st-bc > div:nth-child(5) div.stDownloadButton button[kind='secondary']:hover {
        color: rgb(250, 250, 250);
        background-color: #FF671D;
        background-opacity: 1;
    }


    section.main > div.block-container > div:first-child > div > div.element-container button[title='View fullscreen'] {
        display: none;
    }
    section.main > div.block-container > div:first-child > div > div.element-container img {
        width: 50%;
        margin-left: 30%;
        margin-bottom: 15px;
    }
    div.st-bc > div:nth-child(1) div[data-baseweb='tab-list'] button:first-child {
        margin-left: 33%;
        margin-right: 4%;
    }
    div.st-bc > div:nth-child(1) div[data-baseweb='tab-list'] button:nth-child(2) {
        margin-right: 4%;
    }


    div.st-bc > div:nth-child(3) > div > div > div:nth-child(2) button {
        width: 100%;
    }
    div.st-bc > div:nth-child(3) > div > div > div:nth-child(2) > div:nth-child(3) button {
        border-color: #FF671D;
    }
    div.st-bc > div:nth-child(3) > div > div > div:nth-child(2) > div:nth-child(3) button:hover {
        color: rgb(250, 250, 250);
        background-color: #FF671D;
    }
    div.st-bc > div:nth-child(3) button[title='View fullscreen'] {
        display: none;
    }


    .image_cell {
        width: 160px;
        margin: 10px 20px 0px 0px;
        border-radius: 6px;
        border: 2px solid #0E1117;
        opacity: 0.7;
        display: inline-block;
        position: relative;
    }
    div.st-bc > div:nth-child(4) div[data-testid='stHorizontalBlock'] div[data-testid='stVerticalBlock'] div[data-testid='stVerticalBlock']:hover .image_cell {
        opacity: 1;
        border-color: #FF671D;
        cursor: pointer;
    }
    .image_cell img {
        width:100%;
        border-radius:4px;
    }
    div.st-bc > div:nth-child(4) div[data-testid='stHorizontalBlock'] div[data-testid='stVerticalBlock'] div[data-testid='stVerticalBlock'] {
        margin-bottom: -80px;
    }
    div.st-bc > div:nth-child(4) div[data-testid='stHorizontalBlock'] div[data-testid='stVerticalBlock'] div[data-testid='stVerticalBlock'] div.element-container {
        margin-bottom: -16px;
    }
    div.st-bc > div:nth-child(4) div.stButton button[kind='secondary'] {
        position: relative;
        background-color: rgba(13, 17, 23, 0.4);
        opacity: 0;
        border-radius: 4px;
        padding: 0 10px 0 10px;
    }
    div.st-bc > div:nth-child(4) div[data-testid='stHorizontalBlock'] div[data-testid='stVerticalBlock'] div[data-testid='stVerticalBlock'] div.element-container:nth-child(1) {
        margin-bottom: -24px;
    }
    div.st-bc > div:nth-child(4) div[data-testid='stHorizontalBlock'] div[data-testid='stVerticalBlock'] div[data-testid='stVerticalBlock'] div.element-container:nth-child(2) div.stButton button[kind='secondary'] {
        right: -7px;
        top: -22px;
    }
    div.st-bc > div:nth-child(4) div[data-testid='stHorizontalBlock'] div[data-testid='stVerticalBlock'] div[data-testid='stVerticalBlock'] div.element-container:nth-child(3) div.stButton button[kind='secondary'] {
        right: -57px;
        top: -60px;
        visibility: hidden;
    }
    div.st-bc > div:nth-child(4) div[data-testid='stHorizontalBlock'] div[data-testid='stVerticalBlock'] div[data-testid='stVerticalBlock'] div.element-container:nth-child(4) div.stButton button[kind='secondary'] {
        right: -107px;
        top: -98px;
    }
    div.st-bc > div:nth-child(4) div[data-testid='stHorizontalBlock'] div[data-testid='stVerticalBlock'] div[data-testid='stVerticalBlock']:hover div.stButton button[kind='secondary'] {
        opacity: 1;
    }
    div.st-bc > div:nth-child(4) div.stButton button[kind='secondary']:hover {
        color: rgb(250, 250, 250);
        background-color: #FF671D;
        background-opacity: 1;
    }
    div.st-bc > div:nth-child(4) div.stButton button[kind='secondary'] p {
        font-size: 12px;
    }
    div.st-bc > div:nth-child(4) > div > div > div:nth-child(2) {
        margin-top: 20px;
    }
    div.st-bc > div:nth-child(4) > div > div > div:nth-child(2) > div:nth-child(3) p {
        margin-top: 7px;
    }
    

    div.st-bc > div:nth-child(5) div[data-testid='stExpander']:nth-child(1) div.element-container:nth-child(3),
    div.st-bc > div:nth-child(5) div[data-testid='stExpander']:nth-child(1) div.element-container:nth-child(5) {
        margin-bottom: -16px;
    }

    button[title="View fullscreen"] {
        right: 2px;
        top: 2px;
        color: #FFF;
        background-color: #FF671D;
    }

    footer {
        text-align: center;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

##### header #####

st.image(Image.open(os.path.join("assets", "imgs", "banner.png")), caption="")

tab1, tab2, tab3 = st.tabs(["创作\t:pencil2:", "画廊\t:city_sunrise:", "工具\t:lollipop:"])


##### 创作页面 #####

with tab1:
    prompt_container = st.container()
    col1, col2, col3 = st.columns([4, 1, 1])

    with col1:
        args["negative_prompt"] = st.text_input(
            label="我不想要（可选）",
            value="",
            placeholder="我不想要（可选）",
            label_visibility="collapsed",
        )

    with col2:
        btns["choose_random_prompt"] = st.button("随机描述", on_click=None)
        if btns["choose_random_prompt"]:
            st.session_state.random_prompt = np.random.choice(args["random_prompts"])

    with col3:
        btns["start_generation"] = st.button("创 作", on_click=None)

    with prompt_container:
        args["prompt"] = st.text_area(
            label="描述画面内容",
            value=st.session_state.random_prompt
            if "random_prompt" in st.session_state
            else "",
            placeholder="描述画面内容",
            label_visibility="collapsed",
            height=160,
        )

    with st.container():
        col1, col2 = st.columns([1, 1])
        with col1:
            if os.path.exists(loaded_ckpt_path):
                with open(loaded_ckpt_path, "r") as fr:
                    lines = fr.readlines()

                    if len(lines) > 0:
                        line = lines[0]
                    else:
                        line = args["sd_ckpts"][0] 
            else:
                line = args["sd_ckpts"][0]

            args["model"] = st.selectbox(
                label="选择模型", options=args["sd_ckpts"], index=args["sd_ckpts"].index(line), on_change=switch_ckpt, key="ckpt_name"
            )
            ml = load_ckpt()

            if FREE_HW:
                args["height"] = st.slider(label="高度", min_value=384, max_value=MAX_SIZE, value=512, step=64)
                st.write('\n')

        with col2:
            args["num_images"] = st.number_input(
                label="生成数量", min_value=1, max_value=4, value=1, step=1, on_change=None
            )

            if FREE_HW:
                args["width"] = st.slider(label="宽度", min_value=384, max_value=MAX_SIZE, value=512, step=64)
                st.write('\n')

        if not FREE_HW:
            args["hw"] = st.radio(
                "分辨率 (高 x 宽)",
                ("512 x 512", "768 x 512", "512 x 768", "576 x 576", "640 x 640"),
                index=0,
                horizontal=True,
            )
            args["hw"] = args["hw"].split(" x ")
            args["height"] = int(args["hw"][0])
            args["width"] = int(args["hw"][1])

    col1, col2 = st.columns([1, 2])
    with col1:
        with st.expander("**高级配置**\t:rocket:（可选）", expanded=False):
            args["steps"] = st.slider(
                label="步数（推荐20）", min_value=20, max_value=30, value=20, step=2
            )
            st.write("\n")
            args["cfg"] = st.slider(
                label="引导强度（推荐7.5）", min_value=5.0, max_value=10.0, value=7.5, step=0.5
            )
            st.write("\n")

            if "last_seed" not in st.session_state:
                st.session_state.last_seed = "-1"

            seed_radio = st.radio("随机数种子", ("随机", "上次"), index=0, horizontal=True)
            args["seed"] = st.text_input(
                "种子值（-1表示随机选取）",
                value="-1" if seed_radio == "随机" else st.session_state.last_seed,
                on_change=None,
                placeholder="",
            )
            args["use_enhance"] = st.checkbox("正向增强", value=True)
            args["use_negative_enhance"] = st.checkbox("负向增强", value=True)
            args["use_fast_mode"] = st.checkbox("快速生成", value=True)

    with col2:
        with st.expander("**图像控制**\t:cityscape:（可选）", expanded=False):
            args["init_image"] = st.file_uploader(
                "上传参考图", type=["png", "jpg", "jpeg"], on_change=None
            )

            if args["init_image"] is not None:
                args["init_image"] = np.asarray(
                    bytearray(args["init_image"].read()), dtype=np.uint8
                )
                args["init_image"] = cv2.imdecode(args["init_image"], 1)[:, :, ::-1]
                args["init_image"] = Image.fromarray(
                    args["init_image"].astype("uint8")
                ).convert("RGB")

                w, h = args["init_image"].size
                if max(w, h) > MAX_SIZE:
                    if h > w:
                        args["init_image"] = args["init_image"].resize(
                            (int(MAX_SIZE * w / h), MAX_SIZE)
                        )
                    else:
                        args["init_image"] = args["init_image"].resize(
                            (MAX_SIZE, int(MAX_SIZE * h / w))
                        )

                args["use_init_image"] = True
                st.image(args["init_image"])

                args["denoising_strength"] = st.slider(
                    label="重绘幅度", min_value=0.1, max_value=1.0, value=0.8, step=0.1
                )
                st.write("\n")
            else:
                args["use_init_image"] = False

            args["controlnet_type"] = st.selectbox(
                label="可控生成", options=["不启用", "边缘", "骨骼", "直线"], index=0, on_change=None
            )

            if args["controlnet_type"] != "不启用":
                args["controlnet_image"] = st.file_uploader(
                    "上传控制图", type=["png", "jpg", "jpeg"], on_change=None
                )

                if args["controlnet_image"] is not None:
                    args["controlnet_image"] = np.asarray(
                        bytearray(args["controlnet_image"].read()), dtype=np.uint8
                    )
                    args["controlnet_image"] = cv2.imdecode(
                        args["controlnet_image"], 1
                    )[:, :, ::-1]
                    args["controlnet_image"] = Image.fromarray(
                        args["controlnet_image"].astype("uint8")
                    ).convert("RGB")
                    st.image(args["controlnet_image"])

    with st.expander("**LORA**\t:fire:（可选）", expanded=False):
        st.caption("可用lora，在描述框中使用，可修改权重")
        for lora in args["lora_names"]:
            st.code(lora)

        args['use_dynamic_lora'] = st.checkbox('动态增强', value=True)
        if args['use_dynamic_lora']:
            args['dynamic_mode'] = st.radio('动态增强强度', ('自动', '强', '中', '弱'), index=0, horizontal=True, label_visibility='collapsed')

skip_layers_none = []
skip_layers_some = [
    "lora_unet_mid_block_attentions_0",
    "lora_unet_up_blocks_1_attentions_0",
    "lora_unet_up_blocks_1_attentions_1",
    "lora_unet_up_blocks_1_attentions_2",
]
skip_layers_more = [
    "lora_unet_down_blocks_2_attentions_0",
    "lora_unet_down_blocks_2_attentions_1",
    "lora_unet_mid_block_attentions_0",
    "lora_unet_up_blocks_1_attentions_0",
    "lora_unet_up_blocks_1_attentions_1",
    "lora_unet_up_blocks_1_attentions_2",
    "lora_unet_up_blocks_2_attentions_0",
    "lora_unet_up_blocks_2_attentions_1",
    "lora_unet_up_blocks_2_attentions_2",
]

if args["use_dynamic_lora"]:
    if args["dynamic_mode"] == "自动":
        if args["model"].find("ml-person") >= 0:
            dynamic_skip_layers = skip_layers_more
        elif args["model"].find("ml-2.5D") >= 0:
            dynamic_skip_layers = skip_layers_some
        else:
            dynamic_skip_layers = skip_layers_none
    elif args["dynamic_mode"] == "强":
        dynamic_skip_layers = skip_layers_none
    elif args["dynamic_mode"] == "中":
        dynamic_skip_layers = skip_layers_some
    elif args["dynamic_mode"] == "弱":
        dynamic_skip_layers = skip_layers_more

##### sidebar #####


def progress_callback(step, timestep, image, total_timestep=1000):
    timestep = timestep.to("cpu")
    percentage = (total_timestep - timestep + 1) / total_timestep
    pr_bar.progress(
        int(percentage * 100), text=":smiley:\t生成中，请稍等：{:.2f}%".format(percentage * 100)
    )
    return step, timestep, image


def get_dynamic_lora_from_prompt(prompt, N=3):
    if ml_clip is None:
        return []

    inputs = ml.tokenizer([prompt], padding=True, return_tensors='pt')
    if inputs['input_ids'].shape[1] > 77:
        sos = torch.tensor(49406) # inputs['input_ids'][0][0]
        eos = torch.tensor(49407) # inputs['input_ids'][0][-1]
        inputs['input_ids'] = inputs['input_ids'][:, :77]
        inputs['input_ids'][:, -1] = eos
        inputs['attention_mask'] = inputs['attention_mask'][:, :77]
        inputs['attention_mask'][:, -1] = 0

    text_features = ml_clip.get_text_features(**inputs).data.numpy()

    similarity = F.cosine_similarity(torch.tensor(text_features), torch.tensor(ml_cc), dim=1).data.numpy()
    index = np.argsort(similarity)[::-1]

    idx = index[:N]
    alphas = [similarity[i] for i in idx]

    start, end = 0.2, 0.6
    alphas = [(a - alphas[-1]) / (alphas[0] - alphas[-1]) * (end - start) + start for a in alphas]

    s = np.sum(alphas)
    alphas = [round(r / s, 3) for r in alphas]

    return [[ml_loras_paths[idx[i]], alphas[i]] for i in range(len(idx))]


def load_loras(loras, dynamic=False, skip_layers=[]):
    for lora in loras:
        name, alpha = lora

        if dynamic:
            lora_name = name.split(os.sep)[-1][:-len(".safetensors")]
            lora_path = os.path.join(args["loaded_loras"], "dynamic", f"{lora_name}.pkl")

            if not os.path.exists(lora_path):
                ml.load_lora(
                    name, alpha=alpha, is_path=True, skip_layers=skip_layers
                )

                with open(lora_path, "wb") as fw:
                    pickle.dump([alpha, skip_layers], fw)
        else:
            path = os.path.join("models", "lora", name + ".safetensors")
            lora_path = os.path.join(args["loaded_loras"], "fixed", f"{name}.pkl")
            
            if os.path.exists(path):
                if name not in args["loras"]:
                    args["loras"][name] = load_safetensors(path)

                if not os.path.exists(lora_path):
                    ml.load_lora(
                        args["loras"][name],
                        alpha=alpha,
                        is_path=False,
                        skip_layers=skip_layers,
                    )

                    with open(lora_path, "wb") as fw:
                        pickle.dump(alpha, fw)


def offload_loras(loras, dynamic=False, skip_layers=[]):
    for lora in loras:
        name, alpha = lora

        if dynamic:
            lora_name = name.split(os.sep)[-1][:-len(".safetensors")]
            lora_path = os.path.join(args["loaded_loras"], "dynamic", f"{lora_name}.pkl")

            if os.path.exists(lora_path):
                ml.offload_lora(
                    name, alpha=alpha, is_path=True, skip_layers=skip_layers
                )
                os.system(f"rm {lora_path}")
        else:
            path = os.path.join("models", "lora", name + ".safetensors")
            lora_path = os.path.join(args["loaded_loras"], "fixed", f"{name}.pkl")
            
            if os.path.exists(path):
                if name not in args["loras"]:
                    args["loras"][name] = load_safetensors(path)
                
                if os.path.exists(lora_path):
                    ml.offload_lora(
                        args["loras"][name],
                        alpha=alpha,
                        is_path=False,
                        skip_layers=skip_layers,
                    )
                    os.system(f"rm {lora_path}")


def enhance_prompt(prompt, seed, use_artist=False):
    enhancer = get_enhancer(ml_enhancer[0], seed)
    prompt += ", " + enhancer

    if use_artist:
        artist = get_artists(ml_enhancer[1], seed)
        if artist:
            prompt += ", (" + artist + ":0.6)"

    return prompt


def enhance_negative_prompt(neg):
    safety_neg_prompt, neg_embeddings_prompt, quality_neg_prompt = get_neg_prompt()

    if args["model"].find("ml-general") >= 0:
        neg_embeddings_prompt = ""

    if neg == "":
        neg = safety_neg_prompt + neg_embeddings_prompt + quality_neg_prompt
    else:
        neg = (
            safety_neg_prompt + neg + ", " + neg_embeddings_prompt + quality_neg_prompt
        )
    return neg


with st.sidebar:
    msg = st.container()
    msg.write(" ")

    res = st.container()

    if btns["start_generation"]:
        if args["prompt"] == "":
            msg.write(":slightly_frowning_face:\t请输入画面描述")
        else:
            pr_bar = msg.progress(0, text=":smiley:\t生成中，请稍等")

            if contain_zh(args["prompt"]):
                if ml_translator is None:
                    ml_translator = load_translator()

                args["prompt"] = ml_translator(args["prompt"])

            if contain_zh(args["negative_prompt"]):
                args["negative_prompt"] = ml_translator(args["negative_prompt"])

            if args["seed"] in ["-1", "0"]:
                use_seed = random.randrange(4294967294)
            else:
                try:
                    use_seed = int(args["seed"])
                except:
                    use_seed = random.randrange(4294967294)

            st.session_state.last_seed = str(use_seed)
            st.session_state.last_images = []

            # 加载lora
            fixed_loras, prompt = get_lora_from_prompt(args["prompt"])
            load_loras(fixed_loras)

            if args['use_dynamic_lora']:
                dynamic_loras = get_dynamic_lora_from_prompt(prompt)
                load_loras(dynamic_loras, dynamic=True, skip_layers=dynamic_skip_layers)

            negative_prompt = enhance_negative_prompt(args["negative_prompt"]) if args["use_negative_enhance"] else args["negative_prompt"]

            ori_prompt = prompt
            for i in range(args["num_images"]):
                if args["use_enhance"]:
                    prompt = enhance_prompt(ori_prompt, use_seed + i)

                if args["controlnet_type"] == "不启用" or args["controlnet_image"] is None:
                    if not args["use_init_image"]:
                        images, status = txt2img(
                            sd_model=ml,
                            prompt=prompt,
                            height=args["height"],
                            width=args["width"],
                            num_inference_steps=args["steps"],
                            guidance_scale=args["cfg"],
                            negative_prompt=negative_prompt,
                            num_images_per_prompt=1,
                            seed=use_seed + i,
                            callback=progress_callback,
                            fast_mode=args["use_fast_mode"],
                        )
                    else:
                        images, status = img2img(
                            sd_model=ml,
                            prompt=prompt,
                            image=args["init_image"],
                            strength=args["denoising_strength"],
                            num_inference_steps=args["steps"],
                            guidance_scale=args["cfg"],
                            negative_prompt=negative_prompt,
                            num_images_per_prompt=1,
                            seed=use_seed + i,
                            callback=progress_callback,
                        )

                    anno_image = None

                else:
                    if args["use_init_image"]:
                        w, h = args["init_image"].size
                        if w == args["width"] and h == args["height"]:
                            also_use_init_image = True
                        else:
                            also_use_init_image = False
                    else:
                        also_use_init_image = False

                    if args["controlnet_type"] == "边缘":
                        if mlcn_canny is None:
                            mlcn_canny = load_controlnet("canny")

                        mlcn_canny[1].controlnet = mlcn_canny[1].controlnet.to(device)

                        images, status, anno_image = canny2img(
                            sd_model=ml,
                            controlnet_model=mlcn_canny[1],
                            processor=mlcn_canny[0],
                            ori_image=args["controlnet_image"],
                            aux_image=args["init_image"]
                            if also_use_init_image
                            else None,
                            is_canny_image=False,
                            prompt=prompt,
                            height=args["height"],
                            width=args["width"],
                            strength=1.0,
                            num_inference_steps=args["steps"],
                            guidance_scale=args["cfg"],
                            negative_prompt=negative_prompt,
                            num_images_per_prompt=1,
                            seed=use_seed + i,
                            callback=progress_callback,
                            fast_mode=args["use_fast_mode"],
                        )

                        mlcn_canny[1].controlnet = mlcn_canny[1].controlnet.to("cpu")

                    elif args["controlnet_type"] == "骨骼":
                        if mlcn_pose is None:
                            mlcn_pose = load_controlnet("pose")

                        mlcn_pose[1].controlnet = mlcn_pose[1].controlnet.to(device)

                        images, status, anno_image = pose2img(
                            sd_model=ml,
                            controlnet_model=mlcn_pose[1],
                            processor=mlcn_pose[0],
                            ori_image=args["controlnet_image"],
                            aux_image=args["init_image"]
                            if also_use_init_image
                            else None,
                            is_pose_image=False,
                            prompt=prompt,
                            height=args["height"],
                            width=args["width"],
                            strength=1.0,
                            num_inference_steps=args["steps"],
                            guidance_scale=args["cfg"],
                            negative_prompt=negative_prompt,
                            num_images_per_prompt=1,
                            seed=use_seed + i,
                            callback=progress_callback,
                            fast_mode=args["use_fast_mode"],
                        )

                        mlcn_pose[1].controlnet = mlcn_pose[1].controlnet.to("cpu")

                    elif args["controlnet_type"] == "直线":
                        if mlcn_mlsd is None:
                            mlcn_mlsd = load_controlnet("mlsd")

                        mlcn_mlsd[1].controlnet = mlcn_mlsd[1].controlnet.to(device)

                        images, status, anno_image = mlsd2img(
                            sd_model=ml,
                            controlnet_model=mlcn_mlsd[1],
                            processor=mlcn_mlsd[0],
                            ori_image=args["controlnet_image"],
                            aux_image=args["init_image"]
                            if also_use_init_image
                            else None,
                            is_mlsd_image=False,
                            prompt=prompt,
                            height=args["height"],
                            width=args["width"],
                            strength=1.0,
                            num_inference_steps=args["steps"],
                            guidance_scale=args["cfg"],
                            negative_prompt=negative_prompt,
                            num_images_per_prompt=1,
                            seed=use_seed + i,
                            callback=progress_callback,
                            fast_mode=args["use_fast_mode"],
                        )

                        mlcn_mlsd[1].controlnet = mlcn_mlsd[1].controlnet.to("cpu")

                image = images[0]

                # nsfw
                if status:
                    image = image.filter(
                        ImageFilter.GaussianBlur(
                            radius=max(args["height"], args["width"]) // 8
                        )
                    )

                res.image(image)

                save_name = (
                    str(int(time.time()))
                    + "-"
                    + str(np.random.randint(0, 10e4)).zfill(4)
                    + ".png"
                )

                pnginfo_data = PngImagePlugin.PngInfo()
                value = args["prompt"] + "\n"
                value += "Negative prompt: " + args["negative_prompt"] + "\n"
                value += (
                    "Steps: "
                    + str(args["steps"])
                    + ", Sampler: "
                    + args["sampler"]
                    + ", "
                )
                value += "CFG scale: " + str(args["cfg"]) + ", "
                value += "Seed: " + str(use_seed + i) + ", "
                value += (
                    "Size: " + str(args["width"]) + "x" + str(args["height"]) + ", "
                )
                value += "Model: " + args["model"]
                pnginfo = {"parameters": value}

                for k, v in pnginfo.items():
                    pnginfo_data.add_text(k, str(v))

                save_path = os.path.join(args["save_dir"], save_name)
                image.save(save_path, pnginfo=pnginfo_data)
                st.session_state.last_images.append([image.copy(), save_name])

                w, h = image.size
                image.thumbnail((w // 2, h // 2))
                image.save(
                    os.path.join(args["save_dir"], save_name.rstrip(".png") + ".jpg")
                )

                with open(save_path, "rb") as fr:
                    res.download_button(
                        label="下载", data=fr, file_name=save_name, mime="image/png"
                    )

            if anno_image is not None:
                res.image(anno_image)

            pr_bar.empty()

    elif "last_images" in st.session_state:
        pr_bar = msg.progress(0, text=":smiley:\t生成中，请稍等")
        pr_bar.empty()

        last_images = st.session_state.last_images

        for item in last_images:
            image, save_name = item

            res.image(image)

            save_path = os.path.join(args["save_dir"], save_name)

            if os.path.exists(save_path):
                with open(save_path, "rb") as fr:
                    res.download_button(
                        label="下载", data=fr, file_name=save_name, mime="image/png"
                    )
            else:
                buffered = BytesIO()
                image.save(buffered, format="png")
                res.download_button(
                    label="下载",
                    data=buffered.getvalue(),
                    file_name=save_name,
                    mime="image/png",
                )

##### 画廊页面 #####

with tab2:
    paths = sorted(glob(os.path.join(args["save_dir"], "*.jpg")))[::-1]

    num_image_per_page = 20
    num_of_columns = 4
    total_page = len(paths) // num_image_per_page
    if len(paths) % num_image_per_page > 0:
        total_page += 1

    gallery = st.container()
    col1, col2, col3, col4 = gallery.columns([1, 1, 1, 1])
    columns = [col1, col2, col3, col4]

    if total_page > 1:
        col1, col2, col3 = st.columns([3, 4, 4])
        with col2:
            args["current_page"] = st.number_input(
                label="页码选择",
                min_value=1,
                max_value=total_page,
                value=1,
                step=1,
                on_change=None,
                label_visibility="collapsed",
            )
        with col3:
            st.write(f"共 {total_page} 页")
    else:
        args["current_page"] = 1

    for i, path in enumerate(
        paths[
            (args["current_page"] - 1)
            * num_image_per_page : args["current_page"]
            * num_image_per_page
        ]
    ):
        img = read_img_as_base64(path)

        funcs = columns[i % num_of_columns].container()

        funcs.write(
            f"""
            <div class='image_cell'>
                <img src='{img}'>
            </div>
            """,
            unsafe_allow_html=True,
        )

        name = path.split(os.sep)[-1].rstrip(".jpg")

        tiled_btn = funcs.button("精绘", key=f"hd_{name}", on_click=None)

        if tiled_btn:
            pr_bar = msg.progress(0, text=":smiley:\t精绘中，请稍等")

            st.session_state.last_images = []

            save_name = name + "_tiled.png"
            save_path = os.path.join(args["save_dir"], save_name)

            image = Image.open(os.path.join(args["save_dir"], f"{name}.png"))
            params = read_png_info(image)

            # 加载lora
            fixed_loras, prompt = get_lora_from_prompt(params["prompt"])
            load_loras(fixed_loras)

            if args['use_dynamic_lora']:
                dynamic_loras = get_dynamic_lora_from_prompt(prompt)
                load_loras(dynamic_loras, dynamic=True, skip_layers=dynamic_skip_layers)

            negative_prompt = enhance_negative_prompt(params["negative_prompt"]) if args["use_negative_enhance"] else params["negative_prompt"]

            if args["use_enhance"]:
                prompt = enhance_prompt(prompt, params["seed"])

            if mlcn_tile is None:
                mlcn_tile = load_controlnet("tile")

            mlcn_tile[1].controlnet = mlcn_tile[1].controlnet.to(mlcn_tile[1].device)

            images, status, anno_image = tilehd2img(
                sd_model=ml,
                controlnet_model=mlcn_tile[1],
                processor=mlcn_tile[0],
                ori_image=image,
                is_zoom_image=False,
                up_sampling_ratio=2,
                prompt=prompt,
                # height=params["height"],
                # width=params["width"],
                strength=0.3,
                num_inference_steps=params["steps"],
                guidance_scale=params["cfg"],
                negative_prompt=negative_prompt,
                num_images_per_prompt=1,
                seed=params["seed"],
                callback=progress_callback,
            )

            mlcn_tile[1].controlnet = mlcn_tile[1].controlnet.to("cpu")

            image = images[0]

            # nsfw
            if status:
                image = image.filter(
                    ImageFilter.GaussianBlur(
                        radius=max(params["height"], params["width"]) // 8
                    )
                )

            image.save(save_path)
            st.session_state.last_images.append([image.copy(), save_name])

            res.image(image)

            with open(save_path, 'rb') as fr:
                res.download_button(label='下载', data=fr, key=save_name, file_name=save_name, mime='image/png')

            pr_bar.empty()
            st.experimental_rerun()

        face_btn = funcs.button("修脸", key=f"face_{name}", on_click=None)

        del_btn = funcs.button("删除", key=f"del_{name}", on_click=None)
        if del_btn:
            del_path = os.path.join(args["save_dir"], name)
            os.system(f"rm {del_path}.png")
            os.system(f"rm {del_path}.jpg")

            del_hd_path = f"{del_path}_tiled.png"
            if os.path.exists(del_hd_path):
                os.system(f"rm {del_hd_path}")

            st.experimental_rerun()

##### 工具页面 #####


def record_download_hd():
    st.session_state.hd_image_has_download = True


with tab3:
    with st.expander("**X4超分辨率**\t:rainbow:", expanded=True):
        args["hd_image"] = st.file_uploader(
            "上传图片",
            key="upload_hd_image",
            type=["png", "jpg", "jpeg"],
            on_change=None,
            label_visibility="collapsed",
        )
        if args["hd_image"] is not None:
            filename = args["hd_image"].name
            args["hd_image"] = np.asarray(
                bytearray(args["hd_image"].read()), dtype=np.uint8
            )
            args["hd_image"] = cv2.imdecode(args["hd_image"], 1)[:, :, ::-1]
            args["hd_image"] = Image.fromarray(
                args["hd_image"].astype("uint8")
            ).convert("RGB")

            if (
                "last_hd_image_name" in st.session_state
                and st.session_state.last_hd_image_name == filename
            ):
                hd_image = st.session_state.last_hd_image
            else:
                with st.spinner(text=":smiley:\t超分中，请稍等"):
                    if ml_sr is None:
                        ml_sr = load_sr()

                    hd_image = ml_sr(args["hd_image"])
                    st.session_state.last_hd_image = hd_image
                    st.session_state.last_hd_image_name = filename
                    st.session_state.hd_image_has_download = False

            if (
                "hd_image_has_download" in st.session_state
                and not st.session_state.hd_image_has_download
            ):
                args["hd_image"] = args["hd_image"].resize(hd_image.size)
                combine = np.concatenate(
                    [np.array(args["hd_image"]), np.array(hd_image)], 1
                )
                combine = Image.fromarray(combine.astype("uint8")).convert("RGB")
                st.image(combine)

                buf = BytesIO()
                hd_image.save(buf, format="PNG")
                save_name = "hd_" + filename[: filename.rfind(".")] + ".png"
                st.download_button(
                    label="下载",
                    data=buf.getvalue(),
                    file_name=save_name,
                    mime="image/png",
                    on_click=record_download_hd,
                )

        else:
            st.session_state.last_hd_image = None
            st.session_state.last_hd_image_name = ""
            st.session_state.hd_image_has_download = False

    with st.expander("**获取图片描述**\t:pencil2:", expanded=True):
        args["tag_image"] = st.file_uploader(
            "上传图片",
            key="upload_tag_image",
            type=["png", "jpg", "jpeg"],
            on_change=None,
            label_visibility="collapsed",
        )
        if args["tag_image"] is not None:
            args["tag_image"] = np.asarray(
                bytearray(args["tag_image"].read()), dtype=np.uint8
            )
            args["tag_image"] = cv2.imdecode(args["tag_image"], 1)[:, :, ::-1]
            args["tag_image"] = Image.fromarray(
                args["tag_image"].astype("uint8")
            ).convert("RGB")

            if ml_tagger is None:
                ml_tagger = load_tagger()

            caption = ml_tagger(args["tag_image"])

            st.code(caption)

# 卸载lora
paths = sorted(glob(os.path.join(args["loaded_loras"], "fixed", "*")))
fixed_loras = []
for path in paths:
    with open(path, "rb") as fr:
        alpha = pickle.load(fr)
        fixed_loras.append([path.split(os.sep)[-1][:-len(".pkl")], alpha])
offload_loras(fixed_loras)

paths = sorted(glob(os.path.join(args["loaded_loras"], "dynamic", "*")))
dynamic_loras = []
skip_layers = []
for path in paths:
    with open(path, 'rb') as fr:
        alpha, skip_layers = pickle.load(fr)
        dynamic_loras.append([os.path.join("models", "tools", "dynamic_loras", "lora",
            path.split(os.sep)[-1][:-len('.pkl')] + ".safetensors"), alpha])

offload_loras(dynamic_loras, dynamic=True, skip_layers=skip_layers)
frontend/main.py
MooreThreads/MobiMaliangSDK
def ade_palette():
    """ADE20K palette that maps each class to RGB values."""
    return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
            [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
            [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
            [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
            [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
            [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
            [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
            [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
            [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
            [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
            [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
            [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
            [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
            [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
            [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
            [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
            [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
            [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
            [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
            [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
            [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
            [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
            [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
            [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
            [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
            [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
            [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
            [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
            [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
            [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
            [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
            [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
            [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
            [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
            [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
            [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
            [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
            [102, 255, 0], [92, 0, 255]]

models/controlnet/openpose_v11/controlnet_utils.py
MooreThreads/MobiMaliangSDK
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from diffusers.utils import is_accelerate_available, is_accelerate_version
from typing import Any, Callable, Dict, List, Optional, Union
import gc
import PIL
import numpy as np
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from shared import SDModel, MTGPU_DETECTION


def img2img(
    sd_model: SDModel,
    prompt: Union[str, List[str]] = None,
    image: Union[
        torch.FloatTensor,
        PIL.Image.Image,
        np.ndarray,
        List[torch.FloatTensor],
        List[PIL.Image.Image],
        List[np.ndarray],
    ] = None,
    strength: float = 0.8,
    num_inference_steps: Optional[int] = 20,
    guidance_scale: Optional[float] = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: Optional[float] = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[int] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode: bool = False,
):
    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache() 
    gc.collect()

    img2img_pipeline = StableDiffusionImg2ImgPipeline(
        vae=sd_model.vae,
        text_encoder=sd_model.text_encoder,
        tokenizer=sd_model.tokenizer,
        unet=sd_model.unet,
        scheduler=sd_model.scheduler,
        safety_checker=sd_model.safety_checker,
        feature_extractor=sd_model.feature_extractor,
        requires_safety_checker=sd_model.requires_safety_checker,
    )

    img2img_pipeline.unet.oom = False
    
    if lora_model_name_or_path_or_dict is not None:
        if not isinstance(lora_model_name_or_path_or_dict, list):
            lora_model_name_or_path_or_dict = [lora_model_name_or_path_or_dict]
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.load_lora(lora_weight, alpha=lora_alpha)

    if text_inversion_model_name_or_path is not None:
        if not isinstance(text_inversion_model_name_or_path, list):
            text_inversion_model_name_or_path = [text_inversion_model_name_or_path]
        for text_inversion_weight in text_inversion_model_name_or_path:
            img2img_pipeline.load_textual_inversion(text_inversion_weight)

    generator = None
    if seed is not None:
        if not isinstance(seed, list):
            seed = [seed]
        generator = [
            torch.Generator(sd_model.device).manual_seed(single_seed)
            for single_seed in seed
        ]

    sd_model.unet.oom = False
    
    results = img2img_pipeline(
        prompt=prompt,
        image=image,
        strength=strength,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_images_per_prompt,
        eta=eta,
        generator=generator,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        output_type=output_type,
        return_dict=return_dict,
        callback=callback,
        callback_steps=callback_steps,
        cross_attention_kwargs=cross_attention_kwargs,
    )
    images = results.images
    nsfw_content_detected = results.nsfw_content_detected

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()     
    gc.collect()
    
    if (
        sd_model.enable_offload_cpu
        and is_accelerate_available()
        and is_accelerate_version(">=", "0.14.0")
    ):
        img2img_pipeline.enable_model_cpu_offload()

    if lora_model_name_or_path_or_dict is not None:
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.offload_lora(lora_weight, alpha=lora_alpha)

    return images, nsfw_content_detected

modules/basic/img2img.py
MooreThreads/MobiMaliangSDK
import torch
from diffusers import StableDiffusionPipeline
from diffusers.utils import is_accelerate_available, is_accelerate_version
from typing import Any, Callable, Dict, List, Optional, Union
import gc
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from shared import SDModel, MTGPU_DETECTION


def txt2img(
    sd_model: SDModel,
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    guidance_rescale: float = 0.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode: bool = False,
):

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()    
    gc.collect()

    txt2img_pipeline = StableDiffusionPipeline(
        vae=sd_model.vae,
        text_encoder=sd_model.text_encoder,
        tokenizer=sd_model.tokenizer,
        unet=sd_model.unet,
        scheduler=sd_model.fast_scheduler if fast_mode else sd_model.scheduler,
        safety_checker=sd_model.safety_checker,
        feature_extractor=sd_model.feature_extractor,
        requires_safety_checker=sd_model.requires_safety_checker,
    )
    
    txt2img_pipeline.unet.oom = False

    if lora_model_name_or_path_or_dict is not None:
        if not isinstance(lora_model_name_or_path_or_dict, list):
            lora_model_name_or_path_or_dict = [lora_model_name_or_path_or_dict]
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.load_lora(lora_weight, alpha=lora_alpha)

    if text_inversion_model_name_or_path is not None:
        if not isinstance(text_inversion_model_name_or_path, list):
            text_inversion_model_name_or_path = [text_inversion_model_name_or_path]
        for text_inversion_weight in text_inversion_model_name_or_path:
            txt2img_pipeline.load_textual_inversion(text_inversion_weight)

    generator = None
    if seed is not None:
        if not isinstance(seed, list):
            seed = [seed]
        generator = [
            torch.Generator(sd_model.device).manual_seed(single_seed)
            for single_seed in seed
        ]

    sd_model.unet.oom = False
    
    results = txt2img_pipeline(
        prompt=prompt,
        height=height,
        width=width,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_images_per_prompt,
        eta=eta,
        generator=generator,
        latents=latents,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        output_type=output_type,
        return_dict=return_dict,
        callback=callback,
        callback_steps=callback_steps,
        cross_attention_kwargs=cross_attention_kwargs,
        guidance_rescale=guidance_rescale,
    )
    images = results.images
    nsfw_content_detected = results.nsfw_content_detected
    
    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()    
    gc.collect()

    if (
        sd_model.enable_offload_cpu
        and is_accelerate_available()
        and is_accelerate_version(">=", "0.14.0")
    ):
        txt2img_pipeline.enable_model_cpu_offload()

    if lora_model_name_or_path_or_dict is not None:
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.offload_lora(lora_weight, alpha=lora_alpha)

    return images, nsfw_content_detected

modules/basic/txt2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_txt2img import controlnet_txt2img
from modules.controlnet.controlnet_img2img import controlnet_img2img
from modules.wrappers import CannyDetectorWrapper


def canny2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: CannyDetectorWrapper,
    ori_image: Image,
    aux_image: Image = None,
    is_canny_image: bool = False,
    low_threshold: int = 100,
    high_threshold: int = 200,
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode = False
):

    if is_canny_image:
        canny_image = ori_image
    else:
        canny_image = processor(
            input_image=ori_image,
            low_threshold=low_threshold,
            high_threshold=high_threshold,
        )

    if aux_image is None:
        images, nsfw_content_detected = controlnet_txt2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=canny_image,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
            fast_mode=fast_mode
        )
    else:
        images, nsfw_content_detected = controlnet_img2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=aux_image,
            control_image=canny_image,
            height=height,
            width=width,
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
        )

    return images, nsfw_content_detected, canny_image

modules/controlnet/canny2img.py
MooreThreads/MobiMaliangSDK
import torch
from diffusers import StableDiffusionControlNetImg2ImgPipeline
from diffusers.utils import is_accelerate_available, is_accelerate_version
from typing import Any, Callable, Dict, List, Optional, Union
import gc
import PIL
import numpy as np
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from shared import SDModel, Controlnet, MTGPU_DETECTION


def controlnet_img2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    prompt: Union[str, List[str]] = None,
    image: Union[
        torch.FloatTensor,
        PIL.Image.Image,
        np.ndarray,
        List[torch.FloatTensor],
        List[PIL.Image.Image],
        List[np.ndarray],
    ] = None,
    control_image: Union[
        torch.FloatTensor,
        PIL.Image.Image,
        np.ndarray,
        List[torch.FloatTensor],
        List[PIL.Image.Image],
        List[np.ndarray],
    ] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode: bool = False,
):

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()   
    gc.collect()

    controlnet_img2img_pipeline = StableDiffusionControlNetImg2ImgPipeline(
        vae=sd_model.vae,
        text_encoder=sd_model.text_encoder,
        tokenizer=sd_model.tokenizer,
        unet=sd_model.unet,
        controlnet=controlnet_model.controlnet,
        scheduler=sd_model.scheduler,
        safety_checker=sd_model.safety_checker,
        feature_extractor=sd_model.feature_extractor,
        requires_safety_checker=sd_model.requires_safety_checker,
    )
    
    controlnet_img2img_pipeline.unet.oom = False

    if lora_model_name_or_path_or_dict is not None:
        if not isinstance(lora_model_name_or_path_or_dict, list):
            lora_model_name_or_path_or_dict = [lora_model_name_or_path_or_dict]
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.load_lora(lora_weight, alpha=lora_alpha)

    if text_inversion_model_name_or_path is not None:
        if not isinstance(text_inversion_model_name_or_path, list):
            text_inversion_model_name_or_path = [text_inversion_model_name_or_path]
        for text_inversion_weight in text_inversion_model_name_or_path:
            controlnet_img2img_pipeline.load_textual_inversion(text_inversion_weight)

    generator = None
    if seed is not None:
        if not isinstance(seed, list):
            seed = [seed]
        generator = [
            torch.Generator(sd_model.device).manual_seed(single_seed)
            for single_seed in seed
        ]

    sd_model.unet.oom = False
    
    results = controlnet_img2img_pipeline(
        prompt=prompt,
        image=image,
        control_image=control_image,
        height=height,
        width=width,
        strength=strength,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_images_per_prompt,
        eta=eta,
        generator=generator,
        latents=latents,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        output_type=output_type,
        return_dict=return_dict,
        callback=callback,
        callback_steps=callback_steps,
        cross_attention_kwargs=cross_attention_kwargs,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        guess_mode=guess_mode,
        control_guidance_start=control_guidance_start,
        control_guidance_end=control_guidance_end,
    )
    images = results.images
    nsfw_content_detected = results.nsfw_content_detected

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()    
    gc.collect()
    
    if (
        sd_model.enable_offload_cpu
        and is_accelerate_available()
        and is_accelerate_version(">=", "0.14.0")
    ):
        controlnet_img2img_pipeline.enable_model_cpu_offload()

    if lora_model_name_or_path_or_dict is not None:
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.offload_lora(lora_weight, alpha=lora_alpha)

    return images, nsfw_content_detected

modules/controlnet/controlnet_img2img.py
MooreThreads/MobiMaliangSDK
import torch
from diffusers import StableDiffusionControlNetInpaintPipeline
from diffusers.utils import is_accelerate_available, is_accelerate_version
from typing import Any, Callable, Dict, List, Optional, Union
import gc
import PIL
import numpy as np
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from shared import SDModel, Controlnet, MTGPU_DETECTION


def controlnet_inpaint(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    prompt: Union[str, List[str]] = None,
    image: Union[
        torch.FloatTensor,
        PIL.Image.Image,
        np.ndarray,
        List[torch.FloatTensor],
        List[PIL.Image.Image],
        List[np.ndarray],
    ] = None,
    mask_image: Union[torch.Tensor, PIL.Image.Image] = None,
    control_image: Union[
        torch.FloatTensor,
        PIL.Image.Image,
        np.ndarray,
        List[torch.FloatTensor],
        List[PIL.Image.Image],
        List[np.ndarray],
    ] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode: bool = False,
):

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()    
    gc.collect()

    controlnet_inpaint_pipeline = StableDiffusionControlNetInpaintPipeline(
        vae=sd_model.vae,
        text_encoder=sd_model.text_encoder,
        tokenizer=sd_model.tokenizer,
        unet=sd_model.unet,
        controlnet=controlnet_model.controlnet,
        scheduler=sd_model.scheduler,
        safety_checker=sd_model.safety_checker,
        feature_extractor=sd_model.feature_extractor,
        requires_safety_checker=sd_model.requires_safety_checker,
    )

    controlnet_inpaint_pipeline.unet.oom = False
    
    if lora_model_name_or_path_or_dict is not None:
        if not isinstance(lora_model_name_or_path_or_dict, list):
            lora_model_name_or_path_or_dict = [lora_model_name_or_path_or_dict]
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.load_lora(lora_weight, alpha=lora_alpha)

    if text_inversion_model_name_or_path is not None:
        if not isinstance(text_inversion_model_name_or_path, list):
            text_inversion_model_name_or_path = [text_inversion_model_name_or_path]
        for text_inversion_weight in text_inversion_model_name_or_path:
            controlnet_inpaint_pipeline.load_textual_inversion(text_inversion_weight)

    generator = None
    if seed is not None:
        if not isinstance(seed, list):
            seed = [seed]
        generator = [
            torch.Generator(sd_model.device).manual_seed(single_seed)
            for single_seed in seed
        ]

    sd_model.unet.oom = False
    
    results = controlnet_inpaint_pipeline(
        prompt=prompt,
        image=image,
        mask_image=mask_image,
        control_image=control_image,
        height=height,
        width=width,
        strength=strength,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_images_per_prompt,
        eta=eta,
        generator=generator,
        latents=latents,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        output_type=output_type,
        return_dict=return_dict,
        callback=callback,
        callback_steps=callback_steps,
        cross_attention_kwargs=cross_attention_kwargs,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        guess_mode=guess_mode,
        control_guidance_start=control_guidance_start,
        control_guidance_end=control_guidance_end,
    )
    images = results.images
    nsfw_content_detected = results.nsfw_content_detected

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()    
    gc.collect()
    
    if (
        sd_model.enable_offload_cpu
        and is_accelerate_available()
        and is_accelerate_version(">=", "0.14.0")
    ):
        controlnet_inpaint_pipeline.enable_model_cpu_offload()

    if lora_model_name_or_path_or_dict is not None:
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.offload_lora(lora_weight, alpha=lora_alpha)

    return images, nsfw_content_detected

modules/controlnet/controlnet_inpaint.py
MooreThreads/MobiMaliangSDK
import torch
from diffusers.utils import is_accelerate_available, is_accelerate_version
from typing import Any, Callable, Dict, List, Optional, Union
import gc
import PIL
import numpy as np
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from shared import SDModel, Controlnet, MTGPU_DETECTION
from modules.pipelines.controlnet_tile_hd import (
    StableDiffusionControlNetTileImg2ImgPipeline,
)

# 精绘功能，将图像通过 prompt 控制，进行局部重绘的功能。可以作为加强版的 image to image 功能；也可以用作放大图像功能使用。

def controlnet_tileimg2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    prompt: Union[str, List[str]] = None,
    image: Union[
        torch.FloatTensor,
        PIL.Image.Image,
        np.ndarray,
        List[torch.FloatTensor],
        List[PIL.Image.Image],
        List[np.ndarray],
    ] = None,
    control_image: Union[
        torch.FloatTensor,
        PIL.Image.Image,
        np.ndarray,
        List[torch.FloatTensor],
        List[PIL.Image.Image],
        List[np.ndarray],
    ] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.35,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    window_size: int = 80,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode: bool = False,
):

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()     
    gc.collect()

    tiled_controlnet_img2img_pipeline = StableDiffusionControlNetTileImg2ImgPipeline(
        vae=sd_model.vae,
        text_encoder=sd_model.text_encoder,
        tokenizer=sd_model.tokenizer,
        unet=sd_model.unet.unet,
        controlnet=controlnet_model.controlnet,
        scheduler=sd_model.scheduler,
        safety_checker=sd_model.safety_checker,
        feature_extractor=sd_model.feature_extractor,
        requires_safety_checker=sd_model.requires_safety_checker,
    )
    
    tiled_controlnet_img2img_pipeline.enable_attention_slicing()

    if lora_model_name_or_path_or_dict is not None:
        if not isinstance(lora_model_name_or_path_or_dict, list):
            lora_model_name_or_path_or_dict = [lora_model_name_or_path_or_dict]
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.load_lora(lora_weight, alpha=lora_alpha)

    if text_inversion_model_name_or_path is not None:
        if not isinstance(text_inversion_model_name_or_path, list):
            text_inversion_model_name_or_path = [text_inversion_model_name_or_path]
        for text_inversion_weight in text_inversion_model_name_or_path:
            tiled_controlnet_img2img_pipeline.load_textual_inversion(
                text_inversion_weight
            )

    generator = None
    if seed is not None:
        if not isinstance(seed, list):
            seed = [seed]
        generator = [
            torch.Generator(sd_model.device).manual_seed(single_seed)
            for single_seed in seed
        ]

    sd_model.unet.oom = False
    
    results = tiled_controlnet_img2img_pipeline(
        prompt=prompt,
        image=image,
        control_image=control_image,
        height=height,
        width=width,
        strength=strength,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_images_per_prompt,
        eta=eta,
        generator=generator,
        latents=latents,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        output_type=output_type,
        return_dict=return_dict,
        callback=callback,
        callback_steps=callback_steps,
        cross_attention_kwargs=cross_attention_kwargs,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        guess_mode=guess_mode,
        control_guidance_start=control_guidance_start,
        control_guidance_end=control_guidance_end,
        window_size=window_size,
    )
    images = results.images
    nsfw_content_detected = results.nsfw_content_detected

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()    
    gc.collect()
    
    if (
        sd_model.enable_offload_cpu
        and is_accelerate_available()
        and is_accelerate_version(">=", "0.14.0")
    ):
        tiled_controlnet_img2img_pipeline.enable_model_cpu_offload()

    if lora_model_name_or_path_or_dict is not None:
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.offload_lora(lora_weight, alpha=lora_alpha)

    return images, nsfw_content_detected

modules/controlnet/controlnet_tileimg2img.py
MooreThreads/MobiMaliangSDK
import torch
from diffusers import StableDiffusionControlNetPipeline
from diffusers.utils import is_accelerate_available, is_accelerate_version
from typing import Any, Callable, Dict, List, Optional, Union
import gc
import PIL
import numpy as np
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from shared import SDModel, Controlnet, MTGPU_DETECTION


def controlnet_txt2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    prompt: Union[str, List[str]] = None,
    image: Union[
        torch.FloatTensor,
        PIL.Image.Image,
        np.ndarray,
        List[torch.FloatTensor],
        List[PIL.Image.Image],
        List[np.ndarray],
    ] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode: bool = False,
):

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()    
    gc.collect()

    controlnet_txt2img_pipeline = StableDiffusionControlNetPipeline(
        vae=sd_model.vae,
        text_encoder=sd_model.text_encoder,
        tokenizer=sd_model.tokenizer,
        unet=sd_model.unet,
        controlnet=controlnet_model.controlnet,
        scheduler=sd_model.fast_scheduler if fast_mode else sd_model.scheduler,
        safety_checker=sd_model.safety_checker,
        feature_extractor=sd_model.feature_extractor,
        requires_safety_checker=sd_model.requires_safety_checker,
    )
    
    controlnet_txt2img_pipeline.unet.oom = False

    if lora_model_name_or_path_or_dict is not None:
        if not isinstance(lora_model_name_or_path_or_dict, list):
            lora_model_name_or_path_or_dict = [lora_model_name_or_path_or_dict]
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.load_lora(lora_weight, alpha=lora_alpha)

    if text_inversion_model_name_or_path is not None:
        if not isinstance(text_inversion_model_name_or_path, list):
            text_inversion_model_name_or_path = [text_inversion_model_name_or_path]
        for text_inversion_weight in text_inversion_model_name_or_path:
            controlnet_txt2img_pipeline.load_textual_inversion(text_inversion_weight)

    generator = None
    if seed is not None:
        if not isinstance(seed, list):
            seed = [seed]
        generator = [
            torch.Generator(sd_model.device).manual_seed(single_seed)
            for single_seed in seed
        ]
    
    sd_model.unet.oom = False

    results = controlnet_txt2img_pipeline(
        prompt=prompt,
        image=image,
        height=height,
        width=width,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_images_per_prompt,
        eta=eta,
        generator=generator,
        latents=latents,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        output_type=output_type,
        return_dict=return_dict,
        callback=callback,
        callback_steps=callback_steps,
        cross_attention_kwargs=cross_attention_kwargs,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        guess_mode=guess_mode,
        control_guidance_start=control_guidance_start,
        control_guidance_end=control_guidance_end,
    )
    images = results.images
    nsfw_content_detected = results.nsfw_content_detected

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()    
    gc.collect()
    
    if (
        sd_model.enable_offload_cpu
        and is_accelerate_available()
        and is_accelerate_version(">=", "0.14.0")
    ):
        controlnet_txt2img_pipeline.enable_model_cpu_offload()

    if lora_model_name_or_path_or_dict is not None:
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.offload_lora(lora_weight, alpha=lora_alpha)

    return images, nsfw_content_detected

modules/controlnet/controlnet_txt2img.py
MooreThreads/MobiMaliangSDK
import torch
import numpy as np
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_inpaint import controlnet_inpaint
from modules.wrappers import InpaintingConditionalImageWrapper


def inpaint2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: InpaintingConditionalImageWrapper,
    ori_image: Image,
    mask_image: Image,
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
):

    inpaint_image = processor(input_image=ori_image, mask_image=mask_image)

    images, nsfw_content_detected = controlnet_inpaint(
        sd_model=sd_model,
        controlnet_model=controlnet_model,
        prompt=prompt,
        image=ori_image,
        mask_image=mask_image,
        control_image=inpaint_image,
        height=height,
        width=width,
        strength=strength,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_images_per_prompt,
        eta=eta,
        seed=seed,
        latents=latents,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        output_type=output_type,
        return_dict=return_dict,
        callback=callback,
        callback_steps=callback_steps,
        cross_attention_kwargs=cross_attention_kwargs,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        guess_mode=guess_mode,
        control_guidance_start=control_guidance_start,
        control_guidance_end=control_guidance_end,
        lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
        lora_alpha=lora_alpha,
        text_inversion_model_name_or_path=text_inversion_model_name_or_path,
    )

    inpaint_image = inpaint_image.numpy()[0].transpose(1, 2, 0)
    inpaint_image = np.clip((inpaint_image * 255.0), 0, 255).astype(np.uint8)
    inpaint_image = Image.fromarray(inpaint_image)

    return images, nsfw_content_detected, inpaint_image

modules/controlnet/inpaint2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_txt2img import controlnet_txt2img
from modules.controlnet.controlnet_img2img import controlnet_img2img


def ip2p2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    ori_image: Image,
    aux_image: Image = None,
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode = False
):

    ip2p_image = ori_image

    if aux_image is None:
        images, nsfw_content_detected = controlnet_txt2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=ip2p_image,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
            fast_mode=fast_mode
        )
    else:
        images, nsfw_content_detected = controlnet_img2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=aux_image,
            control_image=ip2p_image,
            height=height,
            width=width,
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
        )

    return images, nsfw_content_detected, ip2p_image

modules/controlnet/ip2p2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_txt2img import controlnet_txt2img
from modules.controlnet.controlnet_img2img import controlnet_img2img
from modules.wrappers import LineartDetectorWrapper


def lineart2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: LineartDetectorWrapper,
    ori_image: Image,
    aux_image: Image = None,
    is_lineart_image: bool = False,
    coarse: bool = False,
    detect_resolution: int = 512,
    image_resolution: int = 512,
    lineart_output_type: str = "pil",
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode = False
):

    if is_lineart_image:
        lineart_image = ori_image
    else:
        lineart_image = processor(
            input_image=ori_image,
            coarse=coarse,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            output_type=lineart_output_type,
        )

    if aux_image is None:
        images, nsfw_content_detected = controlnet_txt2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=lineart_image,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
            fast_mode=fast_mode
        )
    else:
        images, nsfw_content_detected = controlnet_img2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=aux_image,
            control_image=lineart_image,
            height=height,
            width=width,
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
        )

    return images, nsfw_content_detected, lineart_image

modules/controlnet/lineart2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_txt2img import controlnet_txt2img
from modules.controlnet.controlnet_img2img import controlnet_img2img
from modules.wrappers import MLSDdetectorWrapper


def mlsd2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: MLSDdetectorWrapper,
    ori_image: Image,
    aux_image: Image = None,
    is_mlsd_image: bool = False,
    thr_v: float = 0.1,
    thr_d: float = 0.1,
    detect_resolution: int = 512,
    image_resolution: int = 512,
    mlsd_output_type: str = "pil",
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode = False
):

    if is_mlsd_image:
        mlsd_image = ori_image
    else:
        mlsd_image = processor(
            input_image=ori_image,
            thr_v=thr_v,
            thr_d=thr_d,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            output_type=mlsd_output_type,
        )

    if aux_image is None:
        images, nsfw_content_detected = controlnet_txt2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=mlsd_image,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
            fast_mode=fast_mode
        )
    else:
        images, nsfw_content_detected = controlnet_img2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=aux_image,
            control_image=mlsd_image,
            height=height,
            width=width,
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
        )

    return images, nsfw_content_detected, mlsd_image

modules/controlnet/mlsd2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_txt2img import controlnet_txt2img
from modules.controlnet.controlnet_img2img import controlnet_img2img
from modules.wrappers import NormalBaeDetectorWrapper


def normalbae2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: NormalBaeDetectorWrapper,
    ori_image: Image,
    aux_image: Image = None,
    is_normalbae_image: bool = False,
    detect_resolution: int = 512,
    image_resolution: int = 512,
    normalbae_output_type: str = "pil",
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode = False
):

    if is_normalbae_image:
        normalbae_image = ori_image
    else:
        normalbae_image = processor(
            input_image=ori_image,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            output_type=normalbae_output_type,
        )

    if aux_image is None:
        images, nsfw_content_detected = controlnet_txt2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=normalbae_image,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
            fast_mode=fast_mode
        )
    else:
        images, nsfw_content_detected = controlnet_img2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=aux_image,
            control_image=normalbae_image,
            height=height,
            width=width,
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
        )

    return images, nsfw_content_detected, normalbae_image

modules/controlnet/normalbae2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_txt2img import controlnet_txt2img
from modules.controlnet.controlnet_img2img import controlnet_img2img
from modules.wrappers import OpenposeDetectorWrapper


def pose2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: OpenposeDetectorWrapper,
    ori_image: Image,
    aux_image: Image = None,
    is_pose_image: bool = False,
    include_hand: bool = True,
    include_face: bool = True,
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode = False
):

    if is_pose_image:
        pose_image = ori_image
    else:
        pose_image = processor(
            input_image=ori_image, include_hand=include_hand, include_face=include_face
        )

    if aux_image is None:
        images, nsfw_content_detected = controlnet_txt2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=pose_image,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
            fast_mode=fast_mode
        )
    else:
        images, nsfw_content_detected = controlnet_img2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=aux_image,
            control_image=pose_image,
            height=height,
            width=width,
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
        )

    return images, nsfw_content_detected, pose_image

modules/controlnet/pose2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_txt2img import controlnet_txt2img
from modules.controlnet.controlnet_img2img import controlnet_img2img
from modules.wrappers import HEDdetectorWrapper


def scribble2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: HEDdetectorWrapper,
    ori_image: Image,
    aux_image: Image = None,
    is_scribble_image: bool = False,
    detect_resolution: int = 512,
    image_resolution: int = 512,
    safe: bool = False,
    scribble_output_type: str = "pil",
    scribble: bool = False,
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode = False
):

    if is_scribble_image:
        scribble_image = ori_image
    else:
        scribble_image = processor(
            input_image=ori_image,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            output_type=scribble_output_type,
            safe=safe,
            scribble=scribble,
        )

    if aux_image is None:
        images, nsfw_content_detected = controlnet_txt2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=scribble_image,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
            fast_mode=fast_mode
        )
    else:
        images, nsfw_content_detected = controlnet_img2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=aux_image,
            control_image=scribble_image,
            height=height,
            width=width,
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
        )

    return images, nsfw_content_detected, scribble_image

modules/controlnet/scribble2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_txt2img import controlnet_txt2img
from modules.controlnet.controlnet_img2img import controlnet_img2img
from modules.wrappers import ContentShuffleDetectorWrapper


def shuffle2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: ContentShuffleDetectorWrapper,
    ori_image: Image,
    aux_image: Image = None,
    is_shuffle_image: bool = False,
    detect_resolution: int = 512,
    image_resolution: int = 512,
    shuffle_output_type: str = "pil",
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode = False
):

    if is_shuffle_image:
        shuffle_image = ori_image
    else:
        shuffle_image = processor(
            input_image=ori_image,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            output_type=shuffle_output_type,
        )

    if aux_image is None:
        images, nsfw_content_detected = controlnet_txt2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=shuffle_image,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
            fast_mode=fast_mode
        )
    else:
        images, nsfw_content_detected = controlnet_img2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=aux_image,
            control_image=shuffle_image,
            height=height,
            width=width,
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
        )

    return images, nsfw_content_detected, shuffle_image

modules/controlnet/shuffle2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_txt2img import controlnet_txt2img
from modules.controlnet.controlnet_img2img import controlnet_img2img
from modules.wrappers import PidiNetDetectorWrapper


def softedge2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: PidiNetDetectorWrapper,
    ori_image: Image,
    aux_image: Image = None,
    is_softedge_image: bool = False,
    detect_resolution: int = 512,
    image_resolution: int = 512,
    safe: bool = False,
    softedge_output_type: str = "pil",
    scribble: bool = False,
    apply_filter: bool = False,
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.8,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
    fast_mode = False
):

    if is_softedge_image:
        softedge_image = ori_image
    else:
        softedge_image = processor(
            input_image=ori_image,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            output_type=softedge_output_type,
            safe=safe,
            scribble=scribble,
            apply_filter=apply_filter,
        )

    if aux_image is None:
        images, nsfw_content_detected = controlnet_txt2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=softedge_image,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
            fast_mode=fast_mode
        )
    else:
        images, nsfw_content_detected = controlnet_img2img(
            sd_model=sd_model,
            controlnet_model=controlnet_model,
            prompt=prompt,
            image=aux_image,
            control_image=softedge_image,
            height=height,
            width=width,
            strength=strength,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt,
            num_images_per_prompt=num_images_per_prompt,
            eta=eta,
            seed=seed,
            latents=latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            output_type=output_type,
            return_dict=return_dict,
            callback=callback,
            callback_steps=callback_steps,
            cross_attention_kwargs=cross_attention_kwargs,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guess_mode=guess_mode,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
            lora_alpha=lora_alpha,
            text_inversion_model_name_or_path=text_inversion_model_name_or_path,
        )

    return images, nsfw_content_detected, softedge_image

modules/controlnet/softedge2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_img2img import controlnet_img2img
from modules.wrappers import ResizeForConditionalImageWrapper


def tile2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: ResizeForConditionalImageWrapper,
    ori_image: Image,
    is_tile_image: bool = False,
    down_sampling_ratio: float = 8.0,
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 1.0,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
):

    if is_tile_image:
        tile_image = ori_image
    else:
        tile_image = processor(
            input_image=ori_image, down_sampling_ratio=down_sampling_ratio
        )

    images, nsfw_content_detected = controlnet_img2img(
        sd_model=sd_model,
        controlnet_model=controlnet_model,
        prompt=prompt,
        image=tile_image,
        control_image=tile_image,
        height=height,
        width=width,
        strength=strength,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_images_per_prompt,
        eta=eta,
        seed=seed,
        latents=latents,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        output_type=output_type,
        return_dict=return_dict,
        callback=callback,
        callback_steps=callback_steps,
        cross_attention_kwargs=cross_attention_kwargs,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        guess_mode=guess_mode,
        control_guidance_start=control_guidance_start,
        control_guidance_end=control_guidance_end,
        lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
        lora_alpha=lora_alpha,
        text_inversion_model_name_or_path=text_inversion_model_name_or_path,
    )

    return images, nsfw_content_detected, tile_image

modules/controlnet/tile2img.py
MooreThreads/MobiMaliangSDK
import torch
from typing import Any, Callable, Dict, List, Optional, Union
from PIL import Image
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.shared import SDModel, Controlnet
from modules.controlnet.controlnet_tileimg2img import controlnet_tileimg2img
from modules.wrappers import ZoomConditionalImageWrapper


def tilehd2img(
    sd_model: SDModel,
    controlnet_model: Controlnet,
    processor: ZoomConditionalImageWrapper,
    ori_image: Image,
    is_zoom_image: bool = False,
    up_sampling_ratio: float = 2.0,
    prompt: Union[str, List[str]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    strength: float = 0.35,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: float = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
    guess_mode: bool = False,
    control_guidance_start: Union[float, List[float]] = 0.0,
    control_guidance_end: Union[float, List[float]] = 1.0,
    window_size: int = 80,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[float] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
):

    if is_zoom_image:
        zoom_image = ori_image
    else:
        zoom_image = processor(
            input_image=ori_image, up_sampling_ratio=up_sampling_ratio
        )

    images, nsfw_content_detected = controlnet_tileimg2img(
        sd_model=sd_model,
        controlnet_model=controlnet_model,
        prompt=prompt,
        image=zoom_image,
        control_image=zoom_image,
        height=height,
        width=width,
        strength=strength,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_images_per_prompt,
        eta=eta,
        seed=seed,
        latents=latents,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        output_type=output_type,
        return_dict=return_dict,
        callback=callback,
        callback_steps=callback_steps,
        cross_attention_kwargs=cross_attention_kwargs,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        guess_mode=guess_mode,
        control_guidance_start=control_guidance_start,
        control_guidance_end=control_guidance_end,
        window_size=window_size,
        lora_model_name_or_path_or_dict=lora_model_name_or_path_or_dict,
        lora_alpha=lora_alpha,
        text_inversion_model_name_or_path=text_inversion_model_name_or_path,
    )

    return images, nsfw_content_detected, zoom_image

modules/controlnet/tilehd2img.py
MooreThreads/MobiMaliangSDK
from .cpu_offload import (
    replace_sequential_cpu_offload_from_cuda_to_musa,
    replace_model_cpu_offload_from_cuda_to_musa,
)
from .safety_checker_remove_black_image import (
    replace_safety_checker_forward,
    replace_safety_checker_postprocess,
)
from .encode_prompt import replace_encode_prompt
from .convert_ckpt_to_diffusers import replace_convert_controlnet_checkpoint
from .pipeline_stable_diffusion import replace_pipeline_stable_diffusion_call
from .load_from_single_file import replace_from_single_file
modules/hijack/__init__.py
MooreThreads/MobiMaliangSDK
import diffusers
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
    create_unet_diffusers_config,
    convert_ldm_unet_checkpoint,
)
from diffusers import ControlNetModel


def replace_convert_controlnet_checkpoint():
    def convert_controlnet_checkpoint(
        checkpoint,
        original_config,
        checkpoint_path,
        image_size,
        upcast_attention,
        extract_ema,
        use_linear_projection=None,
        cross_attention_dim=None,
    ):
        ctrlnet_config = create_unet_diffusers_config(
            original_config, image_size=image_size, controlnet=True
        )
        ctrlnet_config["upcast_attention"] = upcast_attention

        ctrlnet_config.pop("sample_size")

        if use_linear_projection is not None:
            ctrlnet_config["use_linear_projection"] = use_linear_projection

        if cross_attention_dim is not None:
            ctrlnet_config["cross_attention_dim"] = cross_attention_dim

        # Some controlnet ckpt files are distributed independently from the rest of the
        # model components i.e. https://huggingface.co/thibaud/controlnet-sd21/
        if "time_embed.0.weight" in checkpoint:
            skip_extract_state_dict = True
        else:
            skip_extract_state_dict = False

        converted_ctrl_checkpoint = convert_ldm_unet_checkpoint(
            checkpoint,
            ctrlnet_config,
            path=checkpoint_path,
            extract_ema=extract_ema,
            controlnet=True,
            skip_extract_state_dict=skip_extract_state_dict,
        )
        ctrlnet_config.pop("addition_embed_type")
        ctrlnet_config.pop("addition_time_embed_dim")
        ctrlnet_config.pop("transformer_layers_per_block")
        controlnet_model = ControlNetModel(**ctrlnet_config)

        controlnet_model.load_state_dict(converted_ctrl_checkpoint)

        return controlnet_model

    diffusers.pipelines.stable_diffusion.convert_from_ckpt.convert_controlnet_checkpoint = (
        convert_controlnet_checkpoint
    )

modules/hijack/convert_ckpt_to_diffusers.py
MooreThreads/MobiMaliangSDK
import torch
import diffusers
from diffusers.utils import is_accelerate_available, is_accelerate_version

try:    
    import torch_musa
    MTGPU_DETECTION = True
except:
    print("We cannot import torch_musa")
    MTGPU_DETECTION = False
        

def replace_sequential_cpu_offload_from_cuda_to_musa():
    def enable_sequential_cpu_offload(self, gpu_id=0):
        r"""
        Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
        text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
        `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
        Note that offloading happens on a submodule basis. Memory savings are higher than with
        `enable_model_cpu_offload`, but performance is lower.
        """
        if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
            from accelerate import cpu_offload
        else:
            raise ImportError(
                "`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"
            )
        if MTGPU_DETECTION:
            device = torch.device(f"musa:{gpu_id}")
        else:
            device = torch.device(f"cuda:{gpu_id}") if torch.cuda.is_available() else torch.device("cpu")

        if self.device.type != "cpu":
            self.to("cpu", silence_dtype_warnings=True)
            torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

        for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
            cpu_offload(cpu_offloaded_model, device)

        if self.safety_checker is not None:
            cpu_offload(
                self.safety_checker, execution_device=device, offload_buffers=True
            )

    diffusers.StableDiffusionPipeline.enable_sequential_cpu_offload = (
        enable_sequential_cpu_offload
    )
    diffusers.StableDiffusionImg2ImgPipeline.enable_sequential_cpu_offload = (
        enable_sequential_cpu_offload
    )
    diffusers.StableDiffusionControlNetPipeline.enable_sequential_cpu_offload = (
        enable_sequential_cpu_offload
    )
    diffusers.StableDiffusionControlNetImg2ImgPipeline.enable_sequential_cpu_offload = (
        enable_sequential_cpu_offload
    )
    diffusers.StableDiffusionControlNetInpaintPipeline.enable_sequential_cpu_offload = (
        enable_sequential_cpu_offload
    )


def replace_model_cpu_offload_from_cuda_to_musa():
    def enable_model_cpu_offload(self, gpu_id=0):
        r"""
        Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
        to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
        method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
        `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
        """
        if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
            from accelerate import cpu_offload_with_hook
        else:
            raise ImportError(
                "`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher."
            )

        device = torch.device(f"musa:{gpu_id}")

        if self.device.type != "cpu":
            self.to("cpu", silence_dtype_warnings=True)
            torch.musa.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

        hook = None
        for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
            _, hook = cpu_offload_with_hook(
                cpu_offloaded_model, device, prev_module_hook=hook
            )

        if self.safety_checker is not None:
            _, hook = cpu_offload_with_hook(
                self.safety_checker, device, prev_module_hook=hook
            )

        # We'll offload the last model manually.
        self.final_offload_hook = hook

    diffusers.StableDiffusionPipeline.enable_model_cpu_offload = (
        enable_model_cpu_offload
    )
    diffusers.StableDiffusionImg2ImgPipeline.enable_model_cpu_offload = (
        enable_model_cpu_offload
    )
    diffusers.StableDiffusionControlNetPipeline.enable_model_cpu_offload = (
        enable_model_cpu_offload
    )
    diffusers.StableDiffusionControlNetImg2ImgPipeline.enable_model_cpu_offload = (
        enable_model_cpu_offload
    )
    diffusers.StableDiffusionControlNetInpaintPipeline.enable_model_cpu_offload = (
        enable_model_cpu_offload
    )

modules/hijack/cpu_offload.py
MooreThreads/MobiMaliangSDK
import torch
import diffusers
from diffusers.loaders import (
    FromSingleFileMixin,
    LoraLoaderMixin,
    TextualInversionLoaderMixin,
)
from diffusers.utils import logging
from typing import Any, Callable, Dict, List, Optional, Union
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from modules.tools import prompt_parser

logger = logging.get_logger(__name__)  # pylint: disable=invalid-name


def replace_encode_prompt():
    def _encode_prompt(
        self,
        prompt,
        device,
        num_images_per_prompt,
        do_classifier_free_guidance,
        negative_prompt=None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        lora_scale: Optional[float] = None,
    ):
        r"""
        Encodes the prompt into text encoder hidden states.

        Args:
             prompt (`str` or `List[str]`, *optional*):
                prompt to be encoded
            device: (`torch.device`):
                torch device
            num_images_per_prompt (`int`):
                number of images that should be generated per prompt
            do_classifier_free_guidance (`bool`):
                whether to use classifier free guidance or not
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
            lora_scale (`float`, *optional*):
                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
        """
        # set lora scale so that monkey patched LoRA
        # function of text encoder can correctly access it
        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
            self._lora_scale = lora_scale

        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        if prompt_embeds is None:
            # textual inversion: procecss multi-vector tokens if necessary
            if isinstance(self, TextualInversionLoaderMixin):
                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)

        # get unconditional embeddings for classifier free guidance
        if negative_prompt_embeds is None:
            uncond_tokens: List[str]
            if negative_prompt is None:
                uncond_tokens = [""] * batch_size
            elif prompt is not None and type(prompt) is not type(negative_prompt):
                raise TypeError(
                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
                    f" {type(prompt)}."
                )
            elif isinstance(negative_prompt, str):
                uncond_tokens = [negative_prompt]
            elif batch_size != len(negative_prompt):
                raise ValueError(
                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
                    " the batch size of `prompt`."
                )
            else:
                uncond_tokens = negative_prompt

            # textual inversion: procecss multi-vector tokens if necessary
            if isinstance(self, TextualInversionLoaderMixin):
                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)

        prompt_embeds, negative_prompt_embeds = prompt_parser.get_weighted_text_embeddings(
            pipe=self,
            prompt=prompt,
            uncond_prompt=uncond_tokens,
            max_embeddings_multiples=3,
        )

        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)

        bs_embed, seq_len, _ = prompt_embeds.shape
        # duplicate text embeddings for each generation per prompt, using mps friendly method
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
        prompt_embeds = prompt_embeds.view(
            bs_embed * num_images_per_prompt, seq_len, -1
        )

        if do_classifier_free_guidance:
            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
            seq_len = negative_prompt_embeds.shape[1]

            negative_prompt_embeds = negative_prompt_embeds.to(
                dtype=self.text_encoder.dtype, device=device
            )

            negative_prompt_embeds = negative_prompt_embeds.repeat(
                1, num_images_per_prompt, 1
            )
            negative_prompt_embeds = negative_prompt_embeds.view(
                batch_size * num_images_per_prompt, seq_len, -1
            )

            # For classifier free guidance, we need to do two forward passes.
            # Here we concatenate the unconditional and text embeddings into a single batch
            # to avoid doing two forward passes
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])

        return prompt_embeds

    diffusers.StableDiffusionPipeline._encode_prompt = _encode_prompt
    diffusers.StableDiffusionImg2ImgPipeline._encode_prompt = _encode_prompt
    diffusers.StableDiffusionControlNetPipeline._encode_prompt = _encode_prompt
    diffusers.StableDiffusionControlNetImg2ImgPipeline._encode_prompt = _encode_prompt
    diffusers.StableDiffusionControlNetInpaintPipeline._encode_prompt = _encode_prompt
    try:
        from modules.pipelines.controlnet_tile_hd import (
            StableDiffusionControlNetTileImg2ImgPipeline,
        )

        StableDiffusionControlNetTileImg2ImgPipeline._encode_prompt = _encode_prompt
    except:
        print(
            "can't replace the StableDiffusionControlNetTileImg2ImgPipeline._encode_prompt with default _encode_prompt"
        )

modules/hijack/encode_prompt.py
MooreThreads/MobiMaliangSDK
import diffusers
from pathlib import Path
from diffusers.utils import (
    DIFFUSERS_CACHE,
    HF_HUB_OFFLINE,
    is_safetensors_available,
)

from huggingface_hub import hf_hub_download
def replace_from_single_file():
    @classmethod
    def from_single_file(cls, pretrained_model_link_or_path, **kwargs):
        r"""
        Instantiate a [`DiffusionPipeline`] from pretrained pipeline weights saved in the `.ckpt` format. The pipeline
        is set in evaluation mode (`model.eval()`) by default.

        Parameters:
            pretrained_model_link_or_path (`str` or `os.PathLike`, *optional*):
                Can be either:
                    - A link to the `.ckpt` file (for example
                      `"https://huggingface.co/<repo_id>/blob/main/<path_to_file>.ckpt"`) on the Hub.
                    - A path to a *file* containing all pipeline weights.
            torch_dtype (`str` or `torch.dtype`, *optional*):
                Override the default `torch.dtype` and load the model with another dtype. If `"auto"` is passed, the
                dtype is automatically derived from the model's weights.
            force_download (`bool`, *optional*, defaults to `False`):
                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
                cached versions if they exist.
            cache_dir (`Union[str, os.PathLike]`, *optional*):
                Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
                is not used.
            resume_download (`bool`, *optional*, defaults to `False`):
                Whether or not to resume downloading the model weights and configuration files. If set to `False`, any
                incompletely downloaded files are deleted.
            proxies (`Dict[str, str]`, *optional*):
                A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
            local_files_only (`bool`, *optional*, defaults to `False`):
                Whether to only load local model weights and configuration files or not. If set to True, the model
                won't be downloaded from the Hub.
            use_auth_token (`str` or *bool*, *optional*):
                The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
                `diffusers-cli login` (stored in `~/.huggingface`) is used.
            revision (`str`, *optional*, defaults to `"main"`):
                The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
                allowed by Git.
            use_safetensors (`bool`, *optional*, defaults to `None`):
                If set to `None`, the safetensors weights are downloaded if they're available **and** if the
                safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors
                weights. If set to `False`, safetensors weights are not loaded.
            extract_ema (`bool`, *optional*, defaults to `False`):
                Whether to extract the EMA weights or not. Pass `True` to extract the EMA weights which usually yield
                higher quality images for inference. Non-EMA weights are usually better to continue finetuning.
            upcast_attention (`bool`, *optional*, defaults to `None`):
                Whether the attention computation should always be upcasted.
            image_size (`int`, *optional*, defaults to 512):
                The image size the model was trained on. Use 512 for all Stable Diffusion v1 models and the Stable
                Diffusion v2 base model. Use 768 for Stable Diffusion v2.
            prediction_type (`str`, *optional*):
                The prediction type the model was trained on. Use `'epsilon'` for all Stable Diffusion v1 models and
                the Stable Diffusion v2 base model. Use `'v_prediction'` for Stable Diffusion v2.
            num_in_channels (`int`, *optional*, defaults to `None`):
                The number of input channels. If `None`, it will be automatically inferred.
            scheduler_type (`str`, *optional*, defaults to `"pndm"`):
                Type of scheduler to use. Should be one of `["pndm", "lms", "heun", "euler", "euler-ancestral", "dpm",
                "ddim"]`.
            load_safety_checker (`bool`, *optional*, defaults to `True`):
                Whether to load the safety checker or not.
            text_encoder (`CLIPTextModel`, *optional*, defaults to `None`):
                An instance of
                [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel) to use,
                specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)
                variant. If this parameter is `None`, the function will load a new instance of [CLIP] by itself, if
                needed.
            tokenizer (`CLIPTokenizer`, *optional*, defaults to `None`):
                An instance of
                [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer)
                to use. If this parameter is `None`, the function will load a new instance of [CLIPTokenizer] by
                itself, if needed.
            kwargs (remaining dictionary of keyword arguments, *optional*):
                Can be used to overwrite load and saveable variables (for example the pipeline components of the
                specific pipeline class). The overwritten components are directly passed to the pipelines `__init__`
                method. See example below for more information.

        Examples:

        ```py
        >>> from diffusers import StableDiffusionPipeline

        >>> # Download pipeline from huggingface.co and cache.
        >>> pipeline = StableDiffusionPipeline.from_single_file(
        ...     "https://huggingface.co/WarriorMama777/OrangeMixs/blob/main/Models/AbyssOrangeMix/AbyssOrangeMix.safetensors"
        ... )

        >>> # Download pipeline from local file
        >>> # file is downloaded under ./v1-5-pruned-emaonly.ckpt
        >>> pipeline = StableDiffusionPipeline.from_single_file("./v1-5-pruned-emaonly")

        >>> # Enable float16 and move to GPU
        >>> pipeline = StableDiffusionPipeline.from_single_file(
        ...     "https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.ckpt",
        ...     torch_dtype=torch.float16,
        ... )
        >>> pipeline.to("cuda")
        ```
        """
        # import here to avoid circular dependency
        from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt

        cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
        original_config_file = kwargs.pop("original_config_file", None)
        resume_download = kwargs.pop("resume_download", False)
        force_download = kwargs.pop("force_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE)
        use_auth_token = kwargs.pop("use_auth_token", None)
        revision = kwargs.pop("revision", None)
        extract_ema = kwargs.pop("extract_ema", False)
        image_size = kwargs.pop("image_size", None)
        scheduler_type = kwargs.pop("scheduler_type", "pndm")
        num_in_channels = kwargs.pop("num_in_channels", None)
        upcast_attention = kwargs.pop("upcast_attention", None)
        load_safety_checker = kwargs.pop("load_safety_checker", True)
        prediction_type = kwargs.pop("prediction_type", None)
        text_encoder = kwargs.pop("text_encoder", None)
        tokenizer = kwargs.pop("tokenizer", None)

        torch_dtype = kwargs.pop("torch_dtype", None)

        use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False)

        pipeline_name = cls.__name__
        file_extension = pretrained_model_link_or_path.rsplit(".", 1)[-1]
        from_safetensors = file_extension == "safetensors"

        if from_safetensors and use_safetensors is False:
            raise ValueError("Make sure to install `safetensors` with `pip install safetensors`.")

        # TODO: For now we only support stable diffusion
        stable_unclip = None
        model_type = None
        controlnet = False

        if pipeline_name == "StableDiffusionControlNetPipeline":
            # Model type will be inferred from the checkpoint.
            controlnet = True
        elif "StableDiffusion" in pipeline_name:
            # Model type will be inferred from the checkpoint.
            pass
        elif pipeline_name == "StableUnCLIPPipeline":
            model_type = "FrozenOpenCLIPEmbedder"
            stable_unclip = "txt2img"
        elif pipeline_name == "StableUnCLIPImg2ImgPipeline":
            model_type = "FrozenOpenCLIPEmbedder"
            stable_unclip = "img2img"
        elif pipeline_name == "PaintByExamplePipeline":
            model_type = "PaintByExample"
        elif pipeline_name == "LDMTextToImagePipeline":
            model_type = "LDMTextToImage"
        else:
            raise ValueError(f"Unhandled pipeline class: {pipeline_name}")

        # remove huggingface url
        for prefix in ["https://huggingface.co/", "huggingface.co/", "hf.co/", "https://hf.co/"]:
            if pretrained_model_link_or_path.startswith(prefix):
                pretrained_model_link_or_path = pretrained_model_link_or_path[len(prefix) :]

        # Code based on diffusers.pipelines.pipeline_utils.DiffusionPipeline.from_pretrained
        ckpt_path = Path(pretrained_model_link_or_path)
        if not ckpt_path.is_file():
            # get repo_id and (potentially nested) file path of ckpt in repo
            repo_id = "/".join(ckpt_path.parts[:2])
            file_path = "/".join(ckpt_path.parts[2:])

            if file_path.startswith("blob/"):
                file_path = file_path[len("blob/") :]

            if file_path.startswith("main/"):
                file_path = file_path[len("main/") :]

            pretrained_model_link_or_path = hf_hub_download(
                repo_id,
                filename=file_path,
                cache_dir=cache_dir,
                resume_download=resume_download,
                proxies=proxies,
                local_files_only=local_files_only,
                use_auth_token=use_auth_token,
                revision=revision,
                force_download=force_download,
            )
        pipe = download_from_original_stable_diffusion_ckpt(
            pretrained_model_link_or_path,
            original_config_file=original_config_file,
            pipeline_class=cls,
            model_type=model_type,
            stable_unclip=stable_unclip,
            controlnet=controlnet,
            from_safetensors=from_safetensors,
            extract_ema=extract_ema,
            image_size=image_size,
            scheduler_type=scheduler_type,
            num_in_channels=num_in_channels,
            upcast_attention=upcast_attention,
            load_safety_checker=load_safety_checker,
            prediction_type=prediction_type,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
        )

        if torch_dtype is not None:
            pipe.to(torch_dtype=torch_dtype)

        return pipe
    diffusers.loaders.FromSingleFileMixin.from_single_file = from_single_file

    

modules/hijack/load_from_single_file.py
MooreThreads/MobiMaliangSDK
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import inspect
import warnings
from typing import Any, Callable, Dict, List, Optional, Union

import torch
from packaging import version
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

import diffusers
from diffusers.configuration_utils import FrozenDict
from diffusers.image_processor import VaeImageProcessor
from diffusers.loaders import (
    FromSingleFileMixin,
    LoraLoaderMixin,
    TextualInversionLoaderMixin,
)
from diffusers.models import AutoencoderKL, UNet2DConditionModel
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import (
    deprecate,
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import (
    StableDiffusionSafetyChecker,
)

logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

EXAMPLE_DOC_STRING = """
    Examples:
        ```py
        >>> import torch
        >>> from diffusers import StableDiffusionPipeline

        >>> pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
        >>> pipe = pipe.to("cuda")

        >>> prompt = "a photo of an astronaut riding a horse on mars"
        >>> image = pipe(prompt).images[0]
        ```
"""


def replace_pipeline_stable_diffusion_call():
    @torch.no_grad()
    @replace_example_docstring(EXAMPLE_DOC_STRING)
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        guidance_rescale: float = 0.0,
    ):
        r"""
        Function invoked when calling the pipeline for generation.

        Args:
            prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                instead.
            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The width in pixels of the generated image.
            num_inference_steps (`int`, *optional*, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            guidance_scale (`float`, *optional*, defaults to 7.5):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                `guidance_scale` is defined as `w` of equation 2. of [Imagen
                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
                usually at the expense of lower image quality.
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            eta (`float`, *optional*, defaults to 0.0):
                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
                [`schedulers.DDIMScheduler`], will be ignored for others.
            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
                to make generation deterministic.
            latents (`torch.FloatTensor`, *optional*):
                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor will ge generated by sampling using the supplied random `generator`.
            prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generate image. Choose between
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
                plain tuple.
            callback (`Callable`, *optional*):
                A function that will be called every `callback_steps` steps during inference. The function will be
                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
            callback_steps (`int`, *optional*, defaults to 1):
                The frequency at which the `callback` function will be called. If not specified, the callback will be
                called at every step.
            cross_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
            guidance_rescale (`float`, *optional*, defaults to 0.7):
                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
                Guidance rescale factor should fix overexposure when using zero terminal SNR.

        Examples:

        Returns:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
            When returning a tuple, the first element is a list with the generated images, and the second element is a
            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
            (nsfw) content, according to the `safety_checker`.
        """
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt,
            height,
            width,
            callback_steps,
            negative_prompt,
            prompt_embeds,
            negative_prompt_embeds,
        )

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0

        # 3. Encode input prompt
        text_encoder_lora_scale = (
            cross_attention_kwargs.get("scale", None)
            if cross_attention_kwargs is not None
            else None
        )
        prompt_embeds = self._encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            lora_scale=text_encoder_lora_scale,
        )

        # 4. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        # 5. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )

        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 7. Denoising loop
        will_oom = False
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
                latent_model_input = (
                    torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                )
                latent_model_input = self.scheduler.scale_model_input(
                    latent_model_input, t
                )

                # predict the noise residual
                if will_oom:
                    noise_pred_uncond = self.unet(
                        latent_model_input[:1],
                        t,
                        encoder_hidden_states=prompt_embeds[:1],
                        cross_attention_kwargs=cross_attention_kwargs,
                        return_dict=False,
                    )[0]
                    noise_pred_text = self.unet(
                        latent_model_input[1:],
                        t,
                        encoder_hidden_states=prompt_embeds[1:],
                        cross_attention_kwargs=cross_attention_kwargs,
                        return_dict=False,
                    )[0]
                    noise_pred = torch.cat([noise_pred_uncond, noise_pred_text])
                else:
                    try:
                        noise_pred = self.unet(
                            latent_model_input,
                            t,
                            encoder_hidden_states=prompt_embeds,
                            cross_attention_kwargs=cross_attention_kwargs,
                            return_dict=False,
                        )[0]
                    except:
                        will_oom = True

                        noise_pred_uncond = self.unet(
                            latent_model_input[:1],
                            t,
                            encoder_hidden_states=prompt_embeds[:1],
                            cross_attention_kwargs=cross_attention_kwargs,
                            return_dict=False,
                        )[0]
                        noise_pred_text = self.unet(
                            latent_model_input[1:],
                            t,
                            encoder_hidden_states=prompt_embeds[1:],
                            cross_attention_kwargs=cross_attention_kwargs,
                            return_dict=False,
                        )[0]
                        noise_pred = torch.cat([noise_pred_uncond, noise_pred_text])

                # perform guidance
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + guidance_scale * (
                        noise_pred_text - noise_pred_uncond
                    )

                if do_classifier_free_guidance and guidance_rescale > 0.0:
                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
                    noise_pred = rescale_noise_cfg(
                        noise_pred, noise_pred_text, guidance_rescale=guidance_rescale
                    )

                # compute the previous noisy sample x_t -> x_t-1
                latents = self.scheduler.step(
                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
                )[0]

                # call the callback, if provided
                if i == len(timesteps) - 1 or (
                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
                ):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        callback(i, t, latents)

        if not output_type == "latent":
            image = self.vae.decode(
                latents / self.vae.config.scaling_factor, return_dict=False
            )[0]
            image, has_nsfw_concept = self.run_safety_checker(
                image, device, prompt_embeds.dtype
            )
        else:
            image = latents
            has_nsfw_concept = None

        if has_nsfw_concept is None:
            do_denormalize = [True] * image.shape[0]
        else:
            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]

        image = self.image_processor.postprocess(
            image, output_type=output_type, do_denormalize=do_denormalize
        )

        # Offload last model to CPU
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (image, has_nsfw_concept)

        return StableDiffusionPipelineOutput(
            images=image, nsfw_content_detected=has_nsfw_concept
        )

    diffusers.StableDiffusionPipeline.__call__ = __call__

modules/hijack/pipeline_stable_diffusion.py
MooreThreads/MobiMaliangSDK
import torch
import diffusers
from diffusers.pipelines.stable_diffusion.safety_checker import cosine_distance, logger
from diffusers.utils import deprecate
from typing import Any, Callable, Dict, List, Optional, Union


def replace_safety_checker_forward():
    @torch.no_grad()
    def forward(self, clip_input, images):
        pooled_output = self.vision_model(clip_input)[1]  # pooled_output
        image_embeds = self.visual_projection(pooled_output)

        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
        special_cos_dist = (
            cosine_distance(image_embeds, self.special_care_embeds)
            .cpu()
            .float()
            .numpy()
        )
        cos_dist = (
            cosine_distance(image_embeds, self.concept_embeds).cpu().float().numpy()
        )

        result = []
        batch_size = image_embeds.shape[0]
        for i in range(batch_size):
            result_img = {
                "special_scores": {},
                "special_care": [],
                "concept_scores": {},
                "bad_concepts": [],
            }

            # increase this value to create a stronger `nfsw` filter
            # at the cost of increasing the possibility of filtering benign images
            adjustment = 0.0

            for concept_idx in range(len(special_cos_dist[0])):
                concept_cos = special_cos_dist[i][concept_idx]
                concept_threshold = self.special_care_embeds_weights[concept_idx].item()
                result_img["special_scores"][concept_idx] = round(
                    concept_cos - concept_threshold + adjustment, 3
                )
                if result_img["special_scores"][concept_idx] > 0:
                    result_img["special_care"].append(
                        {concept_idx, result_img["special_scores"][concept_idx]}
                    )
                    adjustment = 0.01

            for concept_idx in range(len(cos_dist[0])):
                concept_cos = cos_dist[i][concept_idx]
                concept_threshold = self.concept_embeds_weights[concept_idx].item()
                result_img["concept_scores"][concept_idx] = round(
                    concept_cos - concept_threshold + adjustment, 3
                )
                if result_img["concept_scores"][concept_idx] > 0:
                    result_img["bad_concepts"].append(concept_idx)

            result.append(result_img)

        has_nsfw_concepts = [len(res["bad_concepts"]) > 0 for res in result]

        if any(has_nsfw_concepts):
            logger.warning(
                "Potential NSFW content was detected in one or more images."
                " Try again with a different prompt and/or seed."
            )

        return images, has_nsfw_concepts

    diffusers.pipelines.stable_diffusion.safety_checker.StableDiffusionSafetyChecker.forward = (
        forward
    )


def replace_safety_checker_postprocess():
    def postprocess(
        self,
        image: torch.FloatTensor,
        output_type: str = "pil",
        do_denormalize: Optional[List[bool]] = None,
    ):
        if not isinstance(image, torch.Tensor):
            raise ValueError(
                f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
            )
        if output_type not in ["latent", "pt", "np", "pil"]:
            deprecation_message = (
                f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "
                "`pil`, `np`, `pt`, `latent`"
            )
            deprecate(
                "Unsupported output_type",
                "1.0.0",
                deprecation_message,
                standard_warn=False,
            )
            output_type = "np"

        if output_type == "latent":
            return image

        image = torch.stack([self.denormalize(image[i]) for i in range(image.shape[0])])

        if output_type == "pt":
            return image

        image = self.pt_to_numpy(image)

        if output_type == "np":
            return image

        if output_type == "pil":
            return self.numpy_to_pil(image)

    diffusers.image_processor.VaeImageProcessor.postprocess = postprocess

modules/hijack/safety_checker_remove_black_image.py
MooreThreads/MobiMaliangSDK
import torch
from diffusers import StableDiffusionControlNetImg2ImgPipeline
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers.models import AutoencoderKL, ControlNetModel, UNet2DConditionModel
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import is_compiled_module, logging, replace_example_docstring

from diffusers.pipelines.stable_diffusion.safety_checker import (
    StableDiffusionSafetyChecker,
)
from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput


from typing import Any, Callable, Dict, List, Optional, Union, Tuple
import PIL
import numpy as np
import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(os.path.dirname(__dir__)))
from hijack import (
    replace_encode_prompt,
    replace_sequential_cpu_offload_from_cuda_to_musa,
)

logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

EXAMPLE_DOC_STRING = """
    Examples:
        ```py
        >>> # !pip install opencv-python transformers accelerate
        >>> from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, UniPCMultistepScheduler
        >>> from diffusers.utils import load_image
        >>> import numpy as np
        >>> import torch

        >>> import cv2
        >>> from PIL import Image

        >>> # download an image
        >>> image = load_image(
        ...     "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png"
        ... )
        >>> input_image = image.convert("RGB")
        >>> w, h = input_image.size
        >>> W, H = int(w * up_sampling_ratio), int(h * up_sampling_ratio)
        >>> H = int(round(H / 8.0)) * 8
        >>> W = int(round(W / 8.0)) * 8
        >>> control_image = input_image.resize((W, H), resample=Image.LANCZOS)

        >>> # load control net and stable diffusion v1-5
        >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-tile", torch_dtype=torch.float32)
        >>> pipe = StableDiffusionControlNetTileImg2ImgPipeline.from_pretrained(
        ...     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float32
        ... )

        >>> # speed up diffusion process with faster scheduler and memory optimization
        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)

        >>> # generate image
        >>> generator = torch.manual_seed(0)
        >>> image = pipe(
        ...     "best quarlity, masterpiece",
        ...     num_inference_steps=20,
        ...     generator=generator,
        ...     image=image,
        ...     control_image=control_image,
        ... ).images[0]
        ```
"""


class StableDiffusionControlNetTileImg2ImgPipeline(
    StableDiffusionControlNetImg2ImgPipeline
):
    def __init__(
        self,
        vae: AutoencoderKL,
        text_encoder: CLIPTextModel,
        tokenizer: CLIPTokenizer,
        unet: UNet2DConditionModel,
        controlnet: Union[
            ControlNetModel,
            List[ControlNetModel],
            Tuple[ControlNetModel],
            MultiControlNetModel,
        ],
        scheduler: KarrasDiffusionSchedulers,
        safety_checker: StableDiffusionSafetyChecker,
        feature_extractor: CLIPImageProcessor,
        requires_safety_checker: bool = True,
    ):
        super().__init__(
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            unet=unet,
            controlnet=controlnet,
            scheduler=scheduler,
            safety_checker=safety_checker,
            feature_extractor=feature_extractor,
            requires_safety_checker=requires_safety_checker,
        )

        # self._encode_prompt = replace_encode_prompt._encode_prompt
        # self.enable_sequential_cpu_offload = replace_sequential_cpu_offload_from_cuda_to_musa.enable_sequential_cpu_offload

    def get_views(self, height, width, window_size=64):
        # Here, we define the mappings F_i (see Eq. 7 in the MultiDiffusion paper https://arxiv.org/abs/2302.08113)
        window_size = int(window_size // 8 * 8)
        stride = window_size - 2
        if height is None or width is None:
            height = self.image_height
            width = self.image_width

        window_height_size = int(min(height / 8, window_size))
        window_width_size = int(min(width / 8, window_size))
        panorama_height = height / 8
        panorama_width = width / 8
        if panorama_height % stride == 0:
            num_blocks_height = int(max(1, panorama_height // stride))
        else:
            num_blocks_height = int(panorama_height // stride + 1)
        if panorama_width % stride == 0:
            num_blocks_width = int(max(1, panorama_width // stride))
        else:
            num_blocks_width = int(panorama_width // stride + 1)
        total_num_blocks = int(num_blocks_height * num_blocks_width)
        views = []
        for i in range(total_num_blocks):
            h_start = int((i // num_blocks_width) * stride)
            h_end = h_start + window_height_size
            if h_end > panorama_height:
                h_start = int(panorama_height - stride)
                h_end = int(panorama_height)
            w_start = int((i % num_blocks_width) * stride)
            w_end = w_start + window_width_size
            if w_end > panorama_width:
                w_start = int(panorama_width - stride)
                w_end = int(panorama_width)
            views.append((h_start, h_end, w_start, w_end))
        return views

    @torch.no_grad()
    @replace_example_docstring(EXAMPLE_DOC_STRING)
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        image: Union[
            torch.FloatTensor,
            PIL.Image.Image,
            np.ndarray,
            List[torch.FloatTensor],
            List[PIL.Image.Image],
            List[np.ndarray],
        ] = None,
        control_image: Union[
            torch.FloatTensor,
            PIL.Image.Image,
            np.ndarray,
            List[torch.FloatTensor],
            List[PIL.Image.Image],
            List[np.ndarray],
        ] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        strength: float = 0.35,
        num_inference_steps: int = 20,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
        guess_mode: bool = False,
        control_guidance_start: Union[float, List[float]] = 0.0,
        control_guidance_end: Union[float, List[float]] = 1.0,
        window_size: int = 80,
    ):
        r"""
        Function invoked when calling the pipeline for generation.

        Args:
            prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                instead.
            image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
                    `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
                The initial image will be used as the starting point for the image generation process. Can also accpet
                image latents as `image`, if passing latents directly, it will not be encoded again.
            control_image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
                    `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
                The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If
                the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. `PIL.Image.Image` can
                also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If
                height and/or width are passed, `image` is resized according to them. If multiple ControlNets are
                specified in init, images must be passed as a list such that each element of the list can be correctly
                batched for input to a single controlnet.
            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The width in pixels of the generated image.
            num_inference_steps (`int`, *optional*, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            guidance_scale (`float`, *optional*, defaults to 7.5):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                `guidance_scale` is defined as `w` of equation 2. of [Imagen
                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
                usually at the expense of lower image quality.
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            eta (`float`, *optional*, defaults to 0.0):
                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
                [`schedulers.DDIMScheduler`], will be ignored for others.
            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
                to make generation deterministic.
            latents (`torch.FloatTensor`, *optional*):
                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor will ge generated by sampling using the supplied random `generator`.
            prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generate image. Choose between
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
                plain tuple.
            callback (`Callable`, *optional*):
                A function that will be called every `callback_steps` steps during inference. The function will be
                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
            callback_steps (`int`, *optional*, defaults to 1):
                The frequency at which the `callback` function will be called. If not specified, the callback will be
                called at every step.
            cross_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
            controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
                The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added
                to the residual in the original unet. If multiple ControlNets are specified in init, you can set the
                corresponding scale as a list. Note that by default, we use a smaller conditioning scale for inpainting
                than for [`~StableDiffusionControlNetPipeline.__call__`].
            guess_mode (`bool`, *optional*, defaults to `False`):
                In this mode, the ControlNet encoder will try best to recognize the content of the input image even if
                you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended.
            control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
                The percentage of total steps at which the controlnet starts applying.
            control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
                The percentage of total steps at which the controlnet stops applying.

        Examples:

        Returns:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
            When returning a tuple, the first element is a list with the generated images, and the second element is a
            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
            (nsfw) content, according to the `safety_checker`.
        """
        controlnet = (
            self.controlnet._orig_mod
            if is_compiled_module(self.controlnet)
            else self.controlnet
        )

        # align format for control guidance
        if not isinstance(control_guidance_start, list) and isinstance(
            control_guidance_end, list
        ):
            control_guidance_start = len(control_guidance_end) * [
                control_guidance_start
            ]
        elif not isinstance(control_guidance_end, list) and isinstance(
            control_guidance_start, list
        ):
            control_guidance_end = len(control_guidance_start) * [control_guidance_end]
        elif not isinstance(control_guidance_start, list) and not isinstance(
            control_guidance_end, list
        ):
            mult = (
                len(controlnet.nets)
                if isinstance(controlnet, MultiControlNetModel)
                else 1
            )
            control_guidance_start, control_guidance_end = (
                mult * [control_guidance_start],
                mult * [control_guidance_end],
            )

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt,
            control_image,
            callback_steps,
            negative_prompt,
            prompt_embeds,
            negative_prompt_embeds,
            controlnet_conditioning_scale,
            control_guidance_start,
            control_guidance_end,
        )

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0

        controlnet = (
            self.controlnet._orig_mod
            if is_compiled_module(self.controlnet)
            else self.controlnet
        )

        if isinstance(controlnet, MultiControlNetModel) and isinstance(
            controlnet_conditioning_scale, float
        ):
            controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(
                controlnet.nets
            )

        global_pool_conditions = (
            controlnet.config.global_pool_conditions
            if isinstance(controlnet, ControlNetModel)
            else controlnet.nets[0].config.global_pool_conditions
        )
        guess_mode = guess_mode or global_pool_conditions

        # 3. Encode input prompt
        text_encoder_lora_scale = (
            cross_attention_kwargs.get("scale", None)
            if cross_attention_kwargs is not None
            else None
        )
        prompt_embeds = self._encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            lora_scale=text_encoder_lora_scale,
        )
        # 4. Prepare image
        image = self.image_processor.preprocess(image).to(dtype=torch.float32)

        # 5. Prepare controlnet_conditioning_image
        if isinstance(controlnet, ControlNetModel):
            control_image = self.prepare_control_image(
                image=control_image,
                width=width,
                height=height,
                batch_size=batch_size * num_images_per_prompt,
                num_images_per_prompt=num_images_per_prompt,
                device=device,
                dtype=controlnet.dtype,
                do_classifier_free_guidance=do_classifier_free_guidance,
                guess_mode=guess_mode,
            )
        elif isinstance(controlnet, MultiControlNetModel):
            control_images = []

            for control_image_ in control_image:
                control_image_ = self.prepare_control_image(
                    image=control_image_,
                    width=width,
                    height=height,
                    batch_size=batch_size * num_images_per_prompt,
                    num_images_per_prompt=num_images_per_prompt,
                    device=device,
                    dtype=controlnet.dtype,
                    do_classifier_free_guidance=do_classifier_free_guidance,
                    guess_mode=guess_mode,
                )

                control_images.append(control_image_)

            control_image = control_images
        else:
            assert False

        # 5. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps, num_inference_steps = self.get_timesteps(
            num_inference_steps, strength, device
        )
        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)

        # 6. Prepare latent variables
        latents = self.prepare_latents(
            image,
            latent_timestep,
            batch_size,
            num_images_per_prompt,
            prompt_embeds.dtype,
            device,
            generator,
        )

        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 7.1 Create tensor stating which controlnets to keep
        controlnet_keep = []
        for i in range(len(timesteps)):
            keeps = [
                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
                for s, e in zip(control_guidance_start, control_guidance_end)
            ]
            controlnet_keep.append(keeps[0] if len(keeps) == 1 else keeps)

        self.image_height = control_image.shape[2]
        self.image_width = control_image.shape[3]

        views = self.get_views(height, width, window_size=window_size)
        count = torch.zeros_like(latents)
        value = torch.zeros_like(latents)

        # 8. Denoising loop
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                count.zero_()
                value.zero_()
                # generate views
                # Here, we iterate through different spatial crops of the latents and denoise them. These
                # denoised (latent) crops are then averaged to produce the final latent
                # for the current timestep via MultiDiffusion. Please see Sec. 4.1 in the
                # MultiDiffusion paper for more details: https://arxiv.org/abs/2302.08113
                for h_start, h_end, w_start, w_end in views:
                    # get the latents corresponding to the current view coordinates
                    latents_for_view = latents[:, :, h_start:h_end, w_start:w_end]
                    controlnet_conditioning_image_for_view = control_image[
                        :, :, h_start * 8 : h_end * 8, w_start * 8 : w_end * 8
                    ]
                    # expand the latents if we are doing classifier free guidance
                    latent_model_input = (
                        torch.cat([latents_for_view] * 2)
                        if do_classifier_free_guidance
                        else latents_for_view
                    )
                    latent_model_input = self.scheduler.scale_model_input(
                        latent_model_input, t
                    )

                    # controlnet(s) inference
                    if guess_mode and do_classifier_free_guidance:
                        # Infer ControlNet only for the conditional batch.
                        control_model_input = latents
                        control_model_input = self.scheduler.scale_model_input(
                            control_model_input, t
                        )
                        controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]
                    else:
                        control_model_input = latent_model_input
                        controlnet_prompt_embeds = prompt_embeds

                    if isinstance(controlnet_keep[i], list):
                        cond_scale = [
                            c * s
                            for c, s in zip(
                                controlnet_conditioning_scale, controlnet_keep[i]
                            )
                        ]
                    else:
                        cond_scale = controlnet_conditioning_scale * controlnet_keep[i]
                    # compute the percentage of total steps we are at
                    current_sampling_percent = i / len(timesteps)

                    if (
                        current_sampling_percent < control_guidance_start[0]
                        or current_sampling_percent > control_guidance_end[0]
                    ):
                        # do not apply the controlnet
                        down_block_res_samples = None
                        mid_block_res_sample = None
                    else:
                        down_block_res_samples, mid_block_res_sample = self.controlnet(
                            control_model_input,
                            t,
                            encoder_hidden_states=controlnet_prompt_embeds,
                            controlnet_cond=controlnet_conditioning_image_for_view,
                            conditioning_scale=cond_scale,
                            guess_mode=guess_mode,
                            return_dict=False,
                        )

                        if guess_mode and do_classifier_free_guidance:
                            # Infered ControlNet only for the conditional batch.
                            # To apply the output of ControlNet to both the unconditional and conditional batches,
                            # add 0 to the unconditional batch to keep it unchanged.
                            down_block_res_samples = [
                                torch.cat([torch.zeros_like(d), d])
                                for d in down_block_res_samples
                            ]
                            mid_block_res_sample = torch.cat(
                                [
                                    torch.zeros_like(mid_block_res_sample),
                                    mid_block_res_sample,
                                ]
                            )

                    # predict the noise residual
                    noise_pred = self.unet(
                        latent_model_input,
                        t,
                        encoder_hidden_states=prompt_embeds,
                        cross_attention_kwargs=cross_attention_kwargs,
                        down_block_additional_residuals=down_block_res_samples,
                        mid_block_additional_residual=mid_block_res_sample,
                        return_dict=False,
                    )[0]

                    # perform guidance
                    if do_classifier_free_guidance:
                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                        noise_pred = noise_pred_uncond + guidance_scale * (
                            noise_pred_text - noise_pred_uncond
                        )

                    # compute the previous noisy sample x_t -> x_t-1
                    latents_view_denoised = self.scheduler.step(
                        noise_pred,
                        t,
                        latents_for_view,
                        **extra_step_kwargs,
                        return_dict=False
                    )[0]

                    value[:, :, h_start:h_end, w_start:w_end] += latents_view_denoised
                    count[:, :, h_start:h_end, w_start:w_end] += 1

                # take the MultiDiffusion step. Eq. 5 in MultiDiffusion paper: https://arxiv.org/abs/2302.08113
                latents = torch.where(count > 0, value / count, value)

                # call the callback, if provided
                if i == len(timesteps) - 1 or (
                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
                ):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        callback(i, t, latents)

        # If we do sequential model offloading, let's offload unet and controlnet
        # manually for max memory savings
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.unet.to("cpu")
            self.controlnet.to("cpu")
            try:
                torch.musa.empty_cache()
            except:
                torch.cuda.empty_cache()

        if not output_type == "latent":
            image = self.vae.decode(
                latents / self.vae.config.scaling_factor, return_dict=False
            )[0]
            image, has_nsfw_concept = self.run_safety_checker(
                image, device, prompt_embeds.dtype
            )
        else:
            image = latents
            has_nsfw_concept = None

        if has_nsfw_concept is None:
            do_denormalize = [True] * image.shape[0]
        else:
            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]

        image = self.image_processor.postprocess(
            image, output_type=output_type, do_denormalize=do_denormalize
        )

        # Offload last model to CPU
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (image, has_nsfw_concept)

        return StableDiffusionPipelineOutput(
            images=image, nsfw_content_detected=has_nsfw_concept
        )

modules/pipelines/controlnet_tile_hd.py
MooreThreads/MobiMaliangSDK
# Inspired by: https://github.com/Mikubill/sd-webui-controlnet/discussions/1236 and https://github.com/Mikubill/sd-webui-controlnet/discussions/1280
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import numpy as np
import PIL.Image
import torch
import gc

from diffusers import StableDiffusionPipeline
from diffusers.models.attention import BasicTransformerBlock
from diffusers.models.unet_2d_blocks import (
    CrossAttnDownBlock2D,
    CrossAttnUpBlock2D,
    DownBlock2D,
    UpBlock2D,
)
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import (
    rescale_noise_cfg,
)
from diffusers.utils import PIL_INTERPOLATION, logging, randn_tensor
from diffusers.utils import is_accelerate_available, is_accelerate_version

import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from shared import SDModel, MTGPU_DETECTION, device

logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

EXAMPLE_DOC_STRING = """
    Examples:
        ```py
        >>> import torch
        >>> from diffusers import UniPCMultistepScheduler
        >>> from diffusers.utils import load_image

        >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")

        >>> pipe = StableDiffusionReferencePipeline.from_pretrained(
                "runwayml/stable-diffusion-v1-5",
                safety_checker=None,
                torch_dtype=torch.float16
                ).to('musa:0')

        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe_controlnet.scheduler.config)

        >>> result_img = pipe(ref_image=input_image,
                        prompt="1girl",
                        num_inference_steps=20,
                        reference_attn=True,
                        reference_adain=True).images[0]

        >>> result_img.show()
        ```
"""


def torch_dfs(model: torch.nn.Module):
    result = [model]
    for child in model.children():
        result += torch_dfs(child)
    return result


class StableDiffusionReferencePipeline(StableDiffusionPipeline):
    def _default_height_width(self, height, width, image):
        # NOTE: It is possible that a list of images have different
        # dimensions for each image, so just checking the first image
        # is not _exactly_ correct, but it is simple.
        while isinstance(image, list):
            image = image[0]

        if height is None:
            if isinstance(image, PIL.Image.Image):
                height = image.height
            elif isinstance(image, torch.Tensor):
                height = image.shape[2]

            height = (height // 8) * 8  # round down to nearest multiple of 8

        if width is None:
            if isinstance(image, PIL.Image.Image):
                width = image.width
            elif isinstance(image, torch.Tensor):
                width = image.shape[3]

            width = (width // 8) * 8  # round down to nearest multiple of 8

        return height, width

    def prepare_image(
        self,
        image,
        width,
        height,
        batch_size,
        num_images_per_prompt,
        device,
        dtype,
        do_classifier_free_guidance=False,
        guess_mode=False,
    ):
        if not isinstance(image, torch.Tensor):
            if isinstance(image, PIL.Image.Image):
                image = [image]

            if isinstance(image[0], PIL.Image.Image):
                images = []

                for image_ in image:
                    image_ = image_.convert("RGB")
                    image_ = image_.resize(
                        (width, height), resample=PIL_INTERPOLATION["lanczos"]
                    )
                    image_ = np.array(image_)
                    image_ = image_[None, :]
                    images.append(image_)

                image = images

                image = np.concatenate(image, axis=0)
                image = np.array(image).astype(np.float32) / 255.0
                image = image * 2 - 1
                image = image.transpose(0, 3, 1, 2)
                image = torch.from_numpy(image)
            elif isinstance(image[0], torch.Tensor):
                image = torch.cat(image, dim=0)

        image_batch_size = image.shape[0]

        if image_batch_size == 1:
            repeat_by = batch_size
        else:
            # image batch size is the same as prompt batch size
            repeat_by = num_images_per_prompt

        image = image.repeat_interleave(repeat_by, dim=0)

        image = image.to(device=device, dtype=dtype)

        if do_classifier_free_guidance and not guess_mode:
            image = torch.cat([image] * 2)

        return image

    def prepare_ref_latents(
        self,
        refimage,
        batch_size,
        dtype,
        device,
        generator,
        do_classifier_free_guidance,
    ):
        refimage = refimage.to(device=device, dtype=dtype)

        # encode the mask image into latents space so we can concatenate it to the latents
        if isinstance(generator, list):
            ref_image_latents = [
                self.vae.encode(refimage[i : i + 1]).latent_dist.sample(
                    generator=generator[i]
                )
                for i in range(batch_size)
            ]
            ref_image_latents = torch.cat(ref_image_latents, dim=0)
        else:
            ref_image_latents = self.vae.encode(refimage).latent_dist.sample(
                generator=generator
            )
        ref_image_latents = self.vae.config.scaling_factor * ref_image_latents

        # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method
        if ref_image_latents.shape[0] < batch_size:
            if not batch_size % ref_image_latents.shape[0] == 0:
                raise ValueError(
                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"
                    f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."
                    " Make sure the number of images that you pass is divisible by the total requested batch size."
                )
            ref_image_latents = ref_image_latents.repeat(
                batch_size // ref_image_latents.shape[0], 1, 1, 1
            )

        ref_image_latents = (
            torch.cat([ref_image_latents] * 2)
            if do_classifier_free_guidance
            else ref_image_latents
        )

        # aligning device to prevent device errors when concating it with the latent model input
        ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)
        return ref_image_latents

    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        ref_image: Union[torch.FloatTensor, PIL.Image.Image] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 20,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        guidance_rescale: float = 0.0,
        attention_auto_machine_weight: float = 1.0,
        gn_auto_machine_weight: float = 1.0,
        style_fidelity: float = 0.5,
        reference_attn: bool = True,
        reference_adain: bool = True,
    ):
        r"""
        Function invoked when calling the pipeline for generation.

        Args:
            prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                instead.
            ref_image (`torch.FloatTensor`, `PIL.Image.Image`):
                The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If
                the type is specified as `Torch.FloatTensor`, it is passed to Reference Control as is. `PIL.Image.Image` can
                also be accepted as an image.
            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The width in pixels of the generated image.
            num_inference_steps (`int`, *optional*, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            guidance_scale (`float`, *optional*, defaults to 7.5):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                `guidance_scale` is defined as `w` of equation 2. of [Imagen
                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
                usually at the expense of lower image quality.
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            eta (`float`, *optional*, defaults to 0.0):
                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
                [`schedulers.DDIMScheduler`], will be ignored for others.
            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
                to make generation deterministic.
            latents (`torch.FloatTensor`, *optional*):
                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor will ge generated by sampling using the supplied random `generator`.
            prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generate image. Choose between
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
                plain tuple.
            callback (`Callable`, *optional*):
                A function that will be called every `callback_steps` steps during inference. The function will be
                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
            callback_steps (`int`, *optional*, defaults to 1):
                The frequency at which the `callback` function will be called. If not specified, the callback will be
                called at every step.
            cross_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
            guidance_rescale (`float`, *optional*, defaults to 0.7):
                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
                Guidance rescale factor should fix overexposure when using zero terminal SNR.
            attention_auto_machine_weight (`float`):
                Weight of using reference query for self attention's context.
                If attention_auto_machine_weight=1.0, use reference query for all self attention's context.
            gn_auto_machine_weight (`float`):
                Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.
            style_fidelity (`float`):
                style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,
                elif style_fidelity=0.0, prompt more important, else balanced.
            reference_attn (`bool`):
                Whether to use reference query for self attention's context.
            reference_adain (`bool`):
                Whether to use reference adain.

        Examples:

        Returns:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
            When returning a tuple, the first element is a list with the generated images, and the second element is a
            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
            (nsfw) content, according to the `safety_checker`.
        """
        assert (
            reference_attn or reference_adain
        ), "`reference_attn` or `reference_adain` must be True."

        # 0. Default height and width to unet
        height, width = self._default_height_width(height, width, ref_image)

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt,
            height,
            width,
            callback_steps,
            negative_prompt,
            prompt_embeds,
            negative_prompt_embeds,
        )

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0

        # 3. Encode input prompt
        text_encoder_lora_scale = (
            cross_attention_kwargs.get("scale", None)
            if cross_attention_kwargs is not None
            else None
        )
        prompt_embeds = self._encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            lora_scale=text_encoder_lora_scale,
        )

        # 4. Preprocess reference image
        ref_image = self.prepare_image(
            image=ref_image,
            width=width,
            height=height,
            batch_size=batch_size * num_images_per_prompt,
            num_images_per_prompt=num_images_per_prompt,
            device=device,
            dtype=prompt_embeds.dtype,
        )

        # 5. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        # 6. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )

        # 7. Prepare reference latent variables
        ref_image_latents = self.prepare_ref_latents(
            ref_image,
            batch_size * num_images_per_prompt,
            prompt_embeds.dtype,
            device,
            generator,
            do_classifier_free_guidance,
        )

        # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 9. Modify self attention and group norm
        MODE = "write"
        uc_mask = (
            torch.Tensor(
                [1] * batch_size * num_images_per_prompt
                + [0] * batch_size * num_images_per_prompt
            )
            .type_as(ref_image_latents)
            .bool()
        )

        def hacked_basic_transformer_inner_forward(
            self,
            hidden_states: torch.FloatTensor,
            attention_mask: Optional[torch.FloatTensor] = None,
            encoder_hidden_states: Optional[torch.FloatTensor] = None,
            encoder_attention_mask: Optional[torch.FloatTensor] = None,
            timestep: Optional[torch.LongTensor] = None,
            cross_attention_kwargs: Dict[str, Any] = None,
            class_labels: Optional[torch.LongTensor] = None,
        ):
            if self.use_ada_layer_norm:
                norm_hidden_states = self.norm1(hidden_states, timestep)
            elif self.use_ada_layer_norm_zero:
                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
                    hidden_states,
                    timestep,
                    class_labels,
                    hidden_dtype=hidden_states.dtype,
                )
            else:
                norm_hidden_states = self.norm1(hidden_states)

            # 1. Self-Attention
            cross_attention_kwargs = (
                cross_attention_kwargs if cross_attention_kwargs is not None else {}
            )
            if self.only_cross_attention:
                attn_output = self.attn1(
                    norm_hidden_states,
                    encoder_hidden_states=encoder_hidden_states
                    if self.only_cross_attention
                    else None,
                    attention_mask=attention_mask,
                    **cross_attention_kwargs,
                )
            else:
                if MODE == "write":
                    self.bank.append(norm_hidden_states.detach().clone())
                    attn_output = self.attn1(
                        norm_hidden_states,
                        encoder_hidden_states=encoder_hidden_states
                        if self.only_cross_attention
                        else None,
                        attention_mask=attention_mask,
                        **cross_attention_kwargs,
                    )
                if MODE == "read":
                    if attention_auto_machine_weight > self.attn_weight:
                        attn_output_uc = self.attn1(
                            norm_hidden_states,
                            encoder_hidden_states=torch.cat(
                                [norm_hidden_states] + self.bank, dim=1
                            ),
                            # attention_mask=attention_mask,
                            **cross_attention_kwargs,
                        )
                        attn_output_c = attn_output_uc.clone()
                        if do_classifier_free_guidance and style_fidelity > 0:
                            attn_output_c[uc_mask] = self.attn1(
                                norm_hidden_states[uc_mask],
                                encoder_hidden_states=norm_hidden_states[uc_mask],
                                **cross_attention_kwargs,
                            )
                        attn_output = (
                            style_fidelity * attn_output_c
                            + (1.0 - style_fidelity) * attn_output_uc
                        )
                        self.bank.clear()
                    else:
                        attn_output = self.attn1(
                            norm_hidden_states,
                            encoder_hidden_states=encoder_hidden_states
                            if self.only_cross_attention
                            else None,
                            attention_mask=attention_mask,
                            **cross_attention_kwargs,
                        )
            if self.use_ada_layer_norm_zero:
                attn_output = gate_msa.unsqueeze(1) * attn_output
            hidden_states = attn_output + hidden_states

            if self.attn2 is not None:
                norm_hidden_states = (
                    self.norm2(hidden_states, timestep)
                    if self.use_ada_layer_norm
                    else self.norm2(hidden_states)
                )

                # 2. Cross-Attention
                attn_output = self.attn2(
                    norm_hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    attention_mask=encoder_attention_mask,
                    **cross_attention_kwargs,
                )
                hidden_states = attn_output + hidden_states

            # 3. Feed-forward
            norm_hidden_states = self.norm3(hidden_states)

            if self.use_ada_layer_norm_zero:
                norm_hidden_states = (
                    norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
                )

            ff_output = self.ff(norm_hidden_states)

            if self.use_ada_layer_norm_zero:
                ff_output = gate_mlp.unsqueeze(1) * ff_output

            hidden_states = ff_output + hidden_states

            return hidden_states

        def hacked_mid_forward(self, *args, **kwargs):
            eps = 1e-6
            x = self.original_forward(*args, **kwargs)
            if MODE == "write":
                if gn_auto_machine_weight >= self.gn_weight:
                    var, mean = torch.var_mean(
                        x.to("cpu"), dim=(2, 3), keepdim=True, correction=0
                    )
                    var = var.to(device)
                    mean = mean.to(device)
                    self.mean_bank.append(mean)
                    self.var_bank.append(var)
            if MODE == "read":
                if len(self.mean_bank) > 0 and len(self.var_bank) > 0:
                    var, mean = torch.var_mean(
                        x.to("cpu"), dim=(2, 3), keepdim=True, correction=0
                    )
                    var = var.to(device)
                    mean = mean.to(device)
                    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5
                    mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))
                    var_acc = sum(self.var_bank) / float(len(self.var_bank))
                    std_acc = (
                        torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5
                    )
                    x_uc = (((x - mean) / std) * std_acc) + mean_acc
                    x_c = x_uc.clone()
                    if do_classifier_free_guidance and style_fidelity > 0:
                        x_c[uc_mask] = x[uc_mask]
                    x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc
                self.mean_bank = []
                self.var_bank = []
            return x

        def hack_CrossAttnDownBlock2D_forward(
            self,
            hidden_states: torch.FloatTensor,
            temb: Optional[torch.FloatTensor] = None,
            encoder_hidden_states: Optional[torch.FloatTensor] = None,
            attention_mask: Optional[torch.FloatTensor] = None,
            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
            encoder_attention_mask: Optional[torch.FloatTensor] = None,
        ):
            eps = 1e-6

            # TODO(Patrick, William) - attention mask is not used
            output_states = ()

            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):
                hidden_states = resnet(hidden_states, temb)
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
                    return_dict=False,
                )[0]
                if MODE == "write":
                    if gn_auto_machine_weight >= self.gn_weight:
                        var, mean = torch.var_mean(
                            hidden_states.to("cpu"),
                            dim=(2, 3),
                            keepdim=True,
                            correction=0,
                        )
                        var = var.to(device)
                        mean = mean.to(device)
                        self.mean_bank.append([mean])
                        self.var_bank.append([var])
                if MODE == "read":
                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:
                        var, mean = torch.var_mean(
                            hidden_states.to("cpu"),
                            dim=(2, 3),
                            keepdim=True,
                            correction=0,
                        )
                        var = var.to(device)
                        mean = mean.to(device)
                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5
                        mean_acc = sum(self.mean_bank[i]) / float(
                            len(self.mean_bank[i])
                        )
                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))
                        std_acc = (
                            torch.maximum(var_acc, torch.zeros_like(var_acc) + eps)
                            ** 0.5
                        )
                        hidden_states_uc = (
                            ((hidden_states - mean) / std) * std_acc
                        ) + mean_acc
                        hidden_states_c = hidden_states_uc.clone()
                        if do_classifier_free_guidance and style_fidelity > 0:
                            hidden_states_c[uc_mask] = hidden_states[uc_mask]
                        hidden_states = (
                            style_fidelity * hidden_states_c
                            + (1.0 - style_fidelity) * hidden_states_uc
                        )

                output_states = output_states + (hidden_states,)

            if MODE == "read":
                self.mean_bank = []
                self.var_bank = []

            if self.downsamplers is not None:
                for downsampler in self.downsamplers:
                    hidden_states = downsampler(hidden_states)

                output_states = output_states + (hidden_states,)

            return hidden_states, output_states

        def hacked_DownBlock2D_forward(self, hidden_states, temb=None):
            eps = 1e-6

            output_states = ()

            for i, resnet in enumerate(self.resnets):
                hidden_states = resnet(hidden_states, temb)

                if MODE == "write":
                    if gn_auto_machine_weight >= self.gn_weight:
                        var, mean = torch.var_mean(
                            hidden_states.to("cpu"),
                            dim=(2, 3),
                            keepdim=True,
                            correction=0,
                        )
                        var = var.to(device)
                        mean = mean.to(device)
                        self.mean_bank.append([mean])
                        self.var_bank.append([var])
                if MODE == "read":
                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:
                        var, mean = torch.var_mean(
                            hidden_states.to("cpu"),
                            dim=(2, 3),
                            keepdim=True,
                            correction=0,
                        )
                        var = var.to(device)
                        mean = mean.to(device)
                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5
                        mean_acc = sum(self.mean_bank[i]) / float(
                            len(self.mean_bank[i])
                        )
                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))
                        std_acc = (
                            torch.maximum(var_acc, torch.zeros_like(var_acc) + eps)
                            ** 0.5
                        )
                        hidden_states_uc = (
                            ((hidden_states - mean) / std) * std_acc
                        ) + mean_acc
                        hidden_states_c = hidden_states_uc.clone()
                        if do_classifier_free_guidance and style_fidelity > 0:
                            hidden_states_c[uc_mask] = hidden_states[uc_mask]
                        hidden_states = (
                            style_fidelity * hidden_states_c
                            + (1.0 - style_fidelity) * hidden_states_uc
                        )

                output_states = output_states + (hidden_states,)

            if MODE == "read":
                self.mean_bank = []
                self.var_bank = []

            if self.downsamplers is not None:
                for downsampler in self.downsamplers:
                    hidden_states = downsampler(hidden_states)

                output_states = output_states + (hidden_states,)

            return hidden_states, output_states

        def hacked_CrossAttnUpBlock2D_forward(
            self,
            hidden_states: torch.FloatTensor,
            res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
            temb: Optional[torch.FloatTensor] = None,
            encoder_hidden_states: Optional[torch.FloatTensor] = None,
            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
            upsample_size: Optional[int] = None,
            attention_mask: Optional[torch.FloatTensor] = None,
            encoder_attention_mask: Optional[torch.FloatTensor] = None,
        ):
            eps = 1e-6
            # TODO(Patrick, William) - attention mask is not used
            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):
                # pop res hidden states
                res_hidden_states = res_hidden_states_tuple[-1]
                res_hidden_states_tuple = res_hidden_states_tuple[:-1]
                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
                hidden_states = resnet(hidden_states, temb)
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
                    return_dict=False,
                )[0]

                if MODE == "write":
                    if gn_auto_machine_weight >= self.gn_weight:
                        var, mean = torch.var_mean(
                            hidden_states.to("cpu"),
                            dim=(2, 3),
                            keepdim=True,
                            correction=0,
                        )
                        var = var.to(device)
                        mean = mean.to(device)
                        self.mean_bank.append([mean])
                        self.var_bank.append([var])
                if MODE == "read":
                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:
                        var, mean = torch.var_mean(
                            hidden_states.to("cpu"),
                            dim=(2, 3),
                            keepdim=True,
                            correction=0,
                        )
                        var = var.to(device)
                        mean = mean.to(device)
                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5
                        mean_acc = sum(self.mean_bank[i]) / float(
                            len(self.mean_bank[i])
                        )
                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))
                        std_acc = (
                            torch.maximum(var_acc, torch.zeros_like(var_acc) + eps)
                            ** 0.5
                        )
                        hidden_states_uc = (
                            ((hidden_states - mean) / std) * std_acc
                        ) + mean_acc
                        hidden_states_c = hidden_states_uc.clone()
                        if do_classifier_free_guidance and style_fidelity > 0:
                            hidden_states_c[uc_mask] = hidden_states[uc_mask]
                        hidden_states = (
                            style_fidelity * hidden_states_c
                            + (1.0 - style_fidelity) * hidden_states_uc
                        )

            if MODE == "read":
                self.mean_bank = []
                self.var_bank = []

            if self.upsamplers is not None:
                for upsampler in self.upsamplers:
                    hidden_states = upsampler(hidden_states, upsample_size)

            return hidden_states

        def hacked_UpBlock2D_forward(
            self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None
        ):
            eps = 1e-6
            for i, resnet in enumerate(self.resnets):
                # pop res hidden states
                res_hidden_states = res_hidden_states_tuple[-1]
                res_hidden_states_tuple = res_hidden_states_tuple[:-1]
                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
                hidden_states = resnet(hidden_states, temb)

                if MODE == "write":
                    if gn_auto_machine_weight >= self.gn_weight:
                        var, mean = torch.var_mean(
                            hidden_states.to("cpu"),
                            dim=(2, 3),
                            keepdim=True,
                            correction=0,
                        )
                        var = var.to(device)
                        mean = mean.to(device)
                        self.mean_bank.append([mean])
                        self.var_bank.append([var])
                if MODE == "read":
                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:
                        var, mean = torch.var_mean(
                            hidden_states.to("cpu"),
                            dim=(2, 3),
                            keepdim=True,
                            correction=0,
                        )
                        var = var.to(device)
                        mean = mean.to(device)
                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5
                        mean_acc = sum(self.mean_bank[i]) / float(
                            len(self.mean_bank[i])
                        )
                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))
                        std_acc = (
                            torch.maximum(var_acc, torch.zeros_like(var_acc) + eps)
                            ** 0.5
                        )
                        hidden_states_uc = (
                            ((hidden_states - mean) / std) * std_acc
                        ) + mean_acc
                        hidden_states_c = hidden_states_uc.clone()
                        if do_classifier_free_guidance and style_fidelity > 0:
                            hidden_states_c[uc_mask] = hidden_states[uc_mask]
                        hidden_states = (
                            style_fidelity * hidden_states_c
                            + (1.0 - style_fidelity) * hidden_states_uc
                        )

            if MODE == "read":
                self.mean_bank = []
                self.var_bank = []

            if self.upsamplers is not None:
                for upsampler in self.upsamplers:
                    hidden_states = upsampler(hidden_states, upsample_size)

            return hidden_states

        if reference_attn:
            attn_modules = [
                module
                for module in torch_dfs(self.unet)
                if isinstance(module, BasicTransformerBlock)
            ]
            attn_modules = sorted(
                attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
            )

            for i, module in enumerate(attn_modules):
                module._original_inner_forward = module.forward
                module.forward = hacked_basic_transformer_inner_forward.__get__(
                    module, BasicTransformerBlock
                )
                module.bank = []
                module.attn_weight = float(i) / float(len(attn_modules))

        if reference_adain:
            gn_modules = [self.unet.mid_block]
            self.unet.mid_block.gn_weight = 0

            down_blocks = self.unet.down_blocks
            for w, module in enumerate(down_blocks):
                module.gn_weight = 1.0 - float(w) / float(len(down_blocks))
                gn_modules.append(module)

            up_blocks = self.unet.up_blocks
            for w, module in enumerate(up_blocks):
                module.gn_weight = float(w) / float(len(up_blocks))
                gn_modules.append(module)

            for i, module in enumerate(gn_modules):
                if getattr(module, "original_forward", None) is None:
                    module.original_forward = module.forward
                if i == 0:
                    # mid_block
                    module.forward = hacked_mid_forward.__get__(module, torch.nn.Module)
                elif isinstance(module, CrossAttnDownBlock2D):
                    module.forward = hack_CrossAttnDownBlock2D_forward.__get__(
                        module, CrossAttnDownBlock2D
                    )
                elif isinstance(module, DownBlock2D):
                    module.forward = hacked_DownBlock2D_forward.__get__(
                        module, DownBlock2D
                    )
                elif isinstance(module, CrossAttnUpBlock2D):
                    module.forward = hacked_CrossAttnUpBlock2D_forward.__get__(
                        module, CrossAttnUpBlock2D
                    )
                elif isinstance(module, UpBlock2D):
                    module.forward = hacked_UpBlock2D_forward.__get__(module, UpBlock2D)
                module.mean_bank = []
                module.var_bank = []
                module.gn_weight *= 2

        # 10. Denoising loop
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
                latent_model_input = (
                    torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                )
                latent_model_input = self.scheduler.scale_model_input(
                    latent_model_input, t
                )

                # ref only part
                noise = randn_tensor(
                    ref_image_latents.shape,
                    generator=generator,
                    device=device,
                    dtype=ref_image_latents.dtype,
                )
                ref_xt = self.scheduler.add_noise(
                    ref_image_latents, noise, t.reshape(1)
                )
                ref_xt = self.scheduler.scale_model_input(ref_xt, t)

                MODE = "write"
                self.unet(
                    ref_xt,
                    t,
                    encoder_hidden_states=prompt_embeds,
                    cross_attention_kwargs=cross_attention_kwargs,
                    return_dict=False,
                )

                # predict the noise residual
                MODE = "read"
                noise_pred = self.unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=prompt_embeds,
                    cross_attention_kwargs=cross_attention_kwargs,
                    return_dict=False,
                )[0]

                # perform guidance
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + guidance_scale * (
                        noise_pred_text - noise_pred_uncond
                    )

                if do_classifier_free_guidance and guidance_rescale > 0.0:
                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
                    noise_pred = rescale_noise_cfg(
                        noise_pred, noise_pred_text, guidance_rescale=guidance_rescale
                    )

                # compute the previous noisy sample x_t -> x_t-1
                latents = self.scheduler.step(
                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
                )[0]

                # call the callback, if provided
                if i == len(timesteps) - 1 or (
                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
                ):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        callback(i, t, latents)

        if not output_type == "latent":
            image = self.vae.decode(
                latents / self.vae.config.scaling_factor, return_dict=False
            )[0]
            image, has_nsfw_concept = self.run_safety_checker(
                image, device, prompt_embeds.dtype
            )
        else:
            image = latents
            has_nsfw_concept = None

        if has_nsfw_concept is None:
            do_denormalize = [True] * image.shape[0]
        else:
            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]

        image = self.image_processor.postprocess(
            image, output_type=output_type, do_denormalize=do_denormalize
        )

        # Offload last model to CPU
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (image, has_nsfw_concept)

        return StableDiffusionPipelineOutput(
            images=image, nsfw_content_detected=has_nsfw_concept
        )


def reference2img(
    sd_model: SDModel,
    prompt: Union[str, List[str]] = None,
    ref_image: Union[torch.FloatTensor, PIL.Image.Image] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    num_inference_steps: Optional[int] = 20,
    guidance_scale: Optional[float] = 7.5,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    eta: Optional[float] = 0.0,
    seed: Optional[Union[int, List[int]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
    callback_steps: int = 1,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    guidance_rescale: float = 0.0,
    attention_auto_machine_weight: float = 1.0,
    gn_auto_machine_weight: float = 1.0,
    style_fidelity: float = 0.5,
    reference_attn: bool = True,
    reference_adain: bool = True,
    lora_model_name_or_path_or_dict: Optional[Union[str, List[str]]] = None,
    lora_alpha: Optional[int] = 0.75,
    text_inversion_model_name_or_path: Optional[
        Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]]
    ] = None,
):

    reference2img_pipeline = StableDiffusionReferencePipeline(
        vae=sd_model.vae,
        text_encoder=sd_model.text_encoder,
        tokenizer=sd_model.tokenizer,
        unet=sd_model.unet,
        scheduler=sd_model.scheduler,
        safety_checker=sd_model.safety_checker,
        feature_extractor=sd_model.feature_extractor,
        requires_safety_checker=sd_model.requires_safety_checker,
    )

    if lora_model_name_or_path_or_dict is not None:
        if not isinstance(lora_model_name_or_path_or_dict, list):
            lora_model_name_or_path_or_dict = [lora_model_name_or_path_or_dict]
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.load_lora(lora_weight, alpha=lora_alpha)

    if text_inversion_model_name_or_path is not None:
        if not isinstance(text_inversion_model_name_or_path, list):
            text_inversion_model_name_or_path = [text_inversion_model_name_or_path]
        for text_inversion_weight in text_inversion_model_name_or_path:
            reference2img_pipeline.load_textual_inversion(text_inversion_weight)

    generator = None
    if seed is not None:
        generator = torch.Generator(sd_model.device).manual_seed(seed)

    results = reference2img_pipeline(
        prompt=prompt,
        ref_image=ref_image,
        height=height,
        width=width,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        num_images_per_prompt=num_images_per_prompt,
        eta=eta,
        generator=generator,
        latents=latents,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        output_type=output_type,
        return_dict=return_dict,
        callback=callback,
        callback_steps=callback_steps,
        cross_attention_kwargs=cross_attention_kwargs,
        guidance_rescale=guidance_rescale,
        attention_auto_machine_weight=attention_auto_machine_weight,
        gn_auto_machine_weight=gn_auto_machine_weight,
        style_fidelity=style_fidelity,
        reference_attn=reference_attn,
        reference_adain=reference_adain,
    )
    images = results.images
    nsfw_content_detected = results.nsfw_content_detected

    if (
        sd_model.enable_offload_cpu
        and is_accelerate_available()
        and is_accelerate_version(">=", "0.14.0")
    ):
        reference2img_pipeline.enable_model_cpu_offload()

    if lora_model_name_or_path_or_dict is not None:
        for lora_weight in lora_model_name_or_path_or_dict:
            sd_model.offload_lora(lora_weight, alpha=lora_alpha)

    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache() 
    gc.collect()

    return images, nsfw_content_detected

modules/reference/reference.py
MooreThreads/MobiMaliangSDK
import torch
from diffusers import StableDiffusionPipeline, ControlNetModel, DDIMScheduler
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
    download_controlnet_from_original_ckpt,
)
from safetensors.torch import load_file
import gc

import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(__dir__)

from modules.tools import anti_oom, samplers

import hijack

hijack.replace_model_cpu_offload_from_cuda_to_musa()
hijack.replace_safety_checker_forward()
hijack.replace_safety_checker_postprocess()
hijack.replace_encode_prompt()
hijack.replace_convert_controlnet_checkpoint()
hijack.replace_from_single_file()

LORA_PREFIX_TEXT_ENCODER = "lora_te"
LORA_PREFIX_UNET = "lora_unet"

try:
    import torch_musa
    MTGPU_DETECTION = True
except:
    MTGPU_DETECTION = False

device = None

if MTGPU_DETECTION: 
    device = "musa"
else:
    device = "cuda" if torch.cuda.is_available() else "cpu"


class SDModel:
    def __init__(self, device, requires_safety_checker=False, enable_offload_cpu=False):

        self.device = device
        self.requires_safety_checker = requires_safety_checker
        self.enable_offload_cpu = enable_offload_cpu

        self.scheduler = None
        self.vae = None
        self.text_encoder = None
        self.tokenizer = None
        self.safety_checker = None
        self.feature_extractor = None
        self.unet = None

    def load_models(self, models_dir_or_single_file, diffusers_format=False):

        if diffusers_format:
            pipe = StableDiffusionPipeline.from_pretrained(
                models_dir_or_single_file, torch_dtype=torch.float32
            ).to(self.device)
        elif (
            os.path.isfile(models_dir_or_single_file)
            and models_dir_or_single_file.endswith(".safetensors")
            or models_dir_or_single_file.endswith(".ckpt")
        ):
            config_file_path = os.path.join(os.path.dirname(models_dir_or_single_file), "v1-inference.yaml")
            try:
                pipe = StableDiffusionPipeline.from_single_file(
                    models_dir_or_single_file,
                    torch_dtype=torch.float32,
                    load_safety_checker=False,
                    original_config_file=config_file_path
                ).to(self.device)
            except:
                pipe = StableDiffusionPipeline.from_single_file(
                    models_dir_or_single_file,
                    torch_dtype=torch.float32,
                    load_safety_checker=False,
                    original_config_file=None
                ).to(self.device)
        else:
            raise ValueError(
                "Make sure the model file ends up with .ckpt or .safetensors or a folder contains model files with diffusers flag which is true"
            )
        self.scheduler = DDIMScheduler(
            num_train_timesteps=pipe.scheduler.config.num_train_timesteps,
            beta_start=pipe.scheduler.config.beta_start,
            beta_end=pipe.scheduler.config.beta_end,
            beta_schedule=pipe.scheduler.config.beta_schedule,
            trained_betas=pipe.scheduler.config.trained_betas,
            clip_sample=pipe.scheduler.config.clip_sample,
            set_alpha_to_one=pipe.scheduler.config.set_alpha_to_one,
            steps_offset=pipe.scheduler.config.steps_offset,
            prediction_type=pipe.scheduler.config.prediction_type,
            thresholding=False,  # pipe.scheduler.config.thresholding,
            dynamic_thresholding_ratio=0.995,  # pipe.scheduler.config.dynamic_thresholding_ratio,
            clip_sample_range=1.0,  # pipe.scheduler.config.clip_sample_range,
            timestep_spacing=pipe.scheduler.config.timestep_spacing,
            rescale_betas_zero_snr=False,  # pipe.scheduler.config.rescale_betas_zero_snr,
        )
        self.fast_scheduler = samplers.BiDDIMScheduler(
            num_train_timesteps=pipe.scheduler.config.num_train_timesteps,
            beta_start=pipe.scheduler.config.beta_start,
            beta_end=pipe.scheduler.config.beta_end,
            beta_schedule=pipe.scheduler.config.beta_schedule,
            trained_betas=pipe.scheduler.config.trained_betas,
            clip_sample=pipe.scheduler.config.clip_sample,
            set_alpha_to_one=pipe.scheduler.config.set_alpha_to_one,
            steps_offset=pipe.scheduler.config.steps_offset,
            prediction_type=pipe.scheduler.config.prediction_type,
            thresholding=False,  # pipe.scheduler.config.thresholding,
            dynamic_thresholding_ratio=0.995,  # pipe.scheduler.config.dynamic_thresholding_ratio,
            clip_sample_range=1.0,  # pipe.scheduler.config.clip_sample_range,
            timestep_spacing=pipe.scheduler.config.timestep_spacing,
            rescale_betas_zero_snr=False,  # pipe.scheduler.config.rescale_betas_zero_snr,
        )

        self.vae = anti_oom.AntiOOMAutoencoderKL(pipe.vae)
        self.text_encoder = pipe.text_encoder
        self.tokenizer = pipe.tokenizer
        if self.requires_safety_checker:
            if pipe.safety_checker is None:
                print("There is no safety checker")
            self.safety_checker = pipe.safety_checker
        self.feature_extractor = pipe.feature_extractor
        self.unet = anti_oom.AntiOOMUNet2DConditionModel(pipe.unet)
        self.device = pipe.device

        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache() 
        gc.collect()

    def offload_models(self):
        self.scheduler = None
        self.vae = None
        self.text_encoder = None
        self.tokenizer = None
        self.safety_checker = None
        self.feature_extractor = None
        self.unet = None
        
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache() 
        gc.collect()

    def load_lora(
        self,
        pretrained_model_name_or_path,
        LORA_PREFIX_UNET=LORA_PREFIX_UNET,
        LORA_PREFIX_TEXT_ENCODER=LORA_PREFIX_TEXT_ENCODER,
        alpha=0.75,
        is_path=True,
        skip_layers=[],
    ):
        assert self.unet is not None
        assert self.text_encoder is not None
        if is_path:
            state_dict = load_file(pretrained_model_name_or_path)
        else:
            state_dict = pretrained_model_name_or_path
        visited = []

        # directly update weight in diffusers model
        for key in state_dict:
            should_skip = False
            for skip_layer in skip_layers:
                if key.find(skip_layer) == 0:
                    should_skip = True
                    break

            if should_skip:
                continue

            # it is suggested to print out the key, it usually will be something like below
            # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"

            # as we have set the alpha beforehand, so just skip
            if ".alpha" in key or key in visited:
                continue

            if "text" in key:
                layer_infos = (
                    key.split(".")[0]
                    .split(LORA_PREFIX_TEXT_ENCODER + "_")[-1]
                    .split("_")
                )
                curr_layer = self.text_encoder

            else:
                layer_infos = (
                    key.split(".")[0].split(LORA_PREFIX_UNET + "_")[-1].split("_")
                )
                
                if isinstance(self.unet, anti_oom.AntiOOMUNet2DConditionModel):
                    curr_layer = self.unet.unet
                else:
                    curr_layer = self.unet

            curr_device, curr_dtype = curr_layer.device, curr_layer.dtype

            # find the target layer
            temp_name = layer_infos.pop(0)
            not_found = False
            while len(layer_infos) > -1:
                try:
                    curr_layer = curr_layer.__getattr__(temp_name)
                    if len(layer_infos) > 0:
                        temp_name = layer_infos.pop(0)
                    elif len(layer_infos) == 0:
                        break
                except Exception:
                    if len(temp_name) > 0:
                        if len(layer_infos) == 0:
                            not_found = True
                            break
                        temp_name += "_" + layer_infos.pop(0)
                    else:
                        temp_name = layer_infos.pop(0)

            if not_found:
                continue

            pair_keys = []
            if "lora_down" in key:
                pair_keys.append(key.replace("lora_down", "lora_up"))
                pair_keys.append(key)
            else:
                pair_keys.append(key)
                pair_keys.append(key.replace("lora_up", "lora_down"))

            # update weight
            if len(state_dict[pair_keys[0]].shape) == 4:
                weight_up = (
                    state_dict[pair_keys[0]]
                    .squeeze(3)
                    .squeeze(2)
                    .to(curr_dtype)
                    .to(curr_device)
                )
                weight_down = (
                    state_dict[pair_keys[1]]
                    .squeeze(3)
                    .squeeze(2)
                    .to(curr_dtype)
                    .to(curr_device)
                )
                curr_layer.weight.data += alpha * torch.mm(
                    weight_up, weight_down
                ).unsqueeze(2).unsqueeze(3)
            else:
                weight_up = state_dict[pair_keys[0]].to(curr_dtype).to(curr_device)
                weight_down = state_dict[pair_keys[1]].to(curr_dtype).to(curr_device)
                curr_layer.weight.data = curr_layer.weight.data + alpha * torch.mm(
                    weight_up, weight_down
                )

            # update visited list
            for item in pair_keys:
                visited.append(item)

    def offload_lora(
        self,
        pretrained_model_name_or_path,
        LORA_PREFIX_UNET=LORA_PREFIX_UNET,
        LORA_PREFIX_TEXT_ENCODER=LORA_PREFIX_TEXT_ENCODER,
        alpha=0.75,
        is_path=True,
        skip_layers=[],
    ):
        assert self.unet is not None
        assert self.text_encoder is not None
        if is_path:
            state_dict = load_file(pretrained_model_name_or_path)
        else:
            state_dict = pretrained_model_name_or_path
        visited = []

        # directly update weight in diffusers model
        for key in state_dict:
            should_skip = False
            for skip_layer in skip_layers:
                if key.find(skip_layer) == 0:
                    should_skip = True
                    break

            if should_skip:
                continue

            # it is suggested to print out the key, it usually will be something like below
            # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"

            # as we have set the alpha beforehand, so just skip
            if ".alpha" in key or key in visited:
                continue

            if "text" in key:
                layer_infos = (
                    key.split(".")[0]
                    .split(LORA_PREFIX_TEXT_ENCODER + "_")[-1]
                    .split("_")
                )
                curr_layer = self.text_encoder
            else:
                layer_infos = (
                    key.split(".")[0].split(LORA_PREFIX_UNET + "_")[-1].split("_")
                )
                if isinstance(self.unet, anti_oom.AntiOOMUNet2DConditionModel):
                    curr_layer = self.unet.unet
                else:
                    curr_layer = self.unet

            curr_device, curr_dtype = curr_layer.device, curr_layer.dtype

            # find the target layer
            temp_name = layer_infos.pop(0)
            not_found = False
            while len(layer_infos) > -1:
                try:
                    curr_layer = curr_layer.__getattr__(temp_name)
                    if len(layer_infos) > 0:
                        temp_name = layer_infos.pop(0)
                    elif len(layer_infos) == 0:
                        break
                except Exception:
                    if len(temp_name) > 0:
                        if len(layer_infos) == 0:
                            not_found = True
                            break
                        temp_name += "_" + layer_infos.pop(0)
                    else:
                        temp_name = layer_infos.pop(0)

            if not_found:
                continue

            pair_keys = []
            if "lora_down" in key:
                pair_keys.append(key.replace("lora_down", "lora_up"))
                pair_keys.append(key)
            else:
                pair_keys.append(key)
                pair_keys.append(key.replace("lora_up", "lora_down"))

            # update weight
            if len(state_dict[pair_keys[0]].shape) == 4:
                weight_up = (
                    state_dict[pair_keys[0]]
                    .squeeze(3)
                    .squeeze(2)
                    .to(curr_dtype)
                    .to(curr_device)
                )
                weight_down = (
                    state_dict[pair_keys[1]]
                    .squeeze(3)
                    .squeeze(2)
                    .to(curr_dtype)
                    .to(curr_device)
                )
                curr_layer.weight.data -= alpha * torch.mm(
                    weight_up, weight_down
                ).unsqueeze(2).unsqueeze(3)
            else:
                weight_up = state_dict[pair_keys[0]].to(curr_dtype).to(curr_device)
                weight_down = state_dict[pair_keys[1]].to(curr_dtype).to(curr_device)
                curr_layer.weight.data -= alpha * torch.mm(weight_up, weight_down)

            # update visited list
            for item in pair_keys:
                visited.append(item)


class Controlnet:
    def __init__(self, device):
        self.device = device
        self.controlnet = None

    def load_models(self, models_dir_or_single_file, diffusers_format=False):

        if diffusers_format:
            self.controlnet = ControlNetModel.from_pretrained(
                models_dir_or_single_file, torch_dtype=torch.float32
            ).to(self.device)
        elif models_dir_or_single_file.endswith(".pth"):
            model_name = os.path.splitext(os.path.split(models_dir_or_single_file)[1])[
                0
            ]
            yaml_file_path = os.path.join(
                os.path.dirname(models_dir_or_single_file), model_name + ".yaml"
            )
            assert (
                yaml_file_path
            ), "Please place the yaml file in the corresponding controlnet path"
            self.controlnet = download_controlnet_from_original_ckpt(
                checkpoint_path=models_dir_or_single_file,
                original_config_file=yaml_file_path,
            ).to(self.device)
        else:
            raise ValueError(
                "Make sure the model file ends up with .pth or a folder contains model files with diffusers flag which is not true"
            )
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache() 
        gc.collect()

    def offload_models(self):
        self.controlnet = None
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache() 
        gc.collect()

modules/shared.py
MooreThreads/MobiMaliangSDK
import torch
from diffusers.models import AutoencoderKL, UNet2DConditionModel
import inspect
from typing import Any, Dict, List, Optional, Tuple, Union
import gc
import os
import sys


try:
    import torch_musa
    MTGPU_DETECTION = True
except:
    MTGPU_DETECTION = False


class AntiOOMUNet2DConditionModel(torch.nn.Module):
    def __init__(self, unet: UNet2DConditionModel):
        super(AntiOOMUNet2DConditionModel, self).__init__()
        self.unet = unet
        self.oom = False
        self.config = unet.config
        self.down_blocks = unet.down_blocks
        self.mid_block = unet.mid_block
        self.up_blocks = unet.up_blocks
        self.device = unet.device
        self.dtype = unet.dtype

    def modules(self):
        return self.unet.modules()

    def parameters(self, recurse: bool = True):
        return self.unet.parameters(recurse=recurse)

    def named_parameters(
        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
    ):
        return self.unet.named_parameters(
            prefix=prefix, recurse=recurse, remove_duplicate=remove_duplicate
        )
        
    def set_attention_slice(
            self, 
            slice_size: Optional[Union[str, int]] = "auto"
        ):
        
        module_names, _ = self._get_signature_keys(self.unet)
        modules = [getattr(self, n, None) for n in module_names]
        modules = [m for m in modules if isinstance(m, torch.nn.Module) and hasattr(m, "set_attention_slice")]

        for module in modules:
            module.set_attention_slice(slice_size)
    
    def _get_signature_keys(self, obj):
        parameters = inspect.signature(obj.__init__).parameters
        required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty}
        optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty})
        expected_modules = set(required_parameters.keys()) - {"self"}
        return expected_modules, optional_parameters

    def forward(
        self,
        sample: torch.FloatTensor,
        timestep: Union[torch.Tensor, float, int],
        encoder_hidden_states: torch.Tensor,
        class_labels: Optional[torch.Tensor] = None,
        timestep_cond: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
        mid_block_additional_residual: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        return_dict: bool = True,
    ):
        if not self.oom:
            try:
                return self.unet(
                    sample=sample,
                    timestep=timestep,
                    encoder_hidden_states=encoder_hidden_states,
                    class_labels=class_labels,
                    timestep_cond=timestep_cond,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
                    added_cond_kwargs=added_cond_kwargs,
                    down_block_additional_residuals=down_block_additional_residuals,
                    mid_block_additional_residual=mid_block_additional_residual,
                    encoder_attention_mask=encoder_attention_mask,
                    return_dict=return_dict,
                )
            except Exception as err:
                if "MUSA out of memory. " in str(err) or "CUDA out of memory. " in str(err):
                    torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache() 
                    gc.collect()
                    self.oom = True
                    # print("We need split the batch into halfs against \"out of memory\"")
                    # self.set_attention_slice()
                    bs = sample.shape[0]
                    half_bs = bs // 2
                    down_block_additional_residuals_half1 = None
                    down_block_additional_residuals_half2 = None
                    if down_block_additional_residuals is not None:
                        down_block_additional_residuals_half1 = [item[:half_bs] for item in down_block_additional_residuals]
                        down_block_additional_residuals_half2 = [item[half_bs:] for item in down_block_additional_residuals]
                    half1 = self.unet(
                        sample=sample[:half_bs],
                        timestep=timestep,
                        encoder_hidden_states=encoder_hidden_states[:half_bs],
                        class_labels=class_labels,
                        timestep_cond=timestep_cond,
                        attention_mask=attention_mask,
                        cross_attention_kwargs=cross_attention_kwargs,
                        added_cond_kwargs=added_cond_kwargs,
                        down_block_additional_residuals=down_block_additional_residuals_half1,
                        mid_block_additional_residual=mid_block_additional_residual[:half_bs],
                        encoder_attention_mask=encoder_attention_mask,
                        return_dict=return_dict,
                    )

                    half2 = self.unet(
                        sample=sample[half_bs:],
                        timestep=timestep,
                        encoder_hidden_states=encoder_hidden_states[half_bs:],
                        class_labels=class_labels,
                        timestep_cond=timestep_cond,
                        attention_mask=attention_mask,
                        cross_attention_kwargs=cross_attention_kwargs,
                        added_cond_kwargs=added_cond_kwargs,
                        down_block_additional_residuals=down_block_additional_residuals_half2,
                        mid_block_additional_residual=mid_block_additional_residual[half_bs:],
                        encoder_attention_mask=encoder_attention_mask,
                        return_dict=return_dict,
                    )
                    return (torch.cat([half1[0], half2[0]]),)
        else:
            bs = sample.shape[0]
            half_bs = bs // 2
            
            down_block_additional_residuals_half1 = None
            down_block_additional_residuals_half2 = None
            if down_block_additional_residuals is not None:
                down_block_additional_residuals_half1 = [item[:half_bs] for item in down_block_additional_residuals]
                down_block_additional_residuals_half2 = [item[half_bs:] for item in down_block_additional_residuals]
            half1 = self.unet(
                sample=sample[:half_bs],
                timestep=timestep,
                encoder_hidden_states=encoder_hidden_states[:half_bs],
                class_labels=class_labels,
                timestep_cond=timestep_cond,
                attention_mask=attention_mask,
                cross_attention_kwargs=cross_attention_kwargs,
                added_cond_kwargs=added_cond_kwargs,
                down_block_additional_residuals=down_block_additional_residuals_half1,
                mid_block_additional_residual=mid_block_additional_residual[:half_bs],
                encoder_attention_mask=encoder_attention_mask,
                return_dict=return_dict,
            )

            half2 = self.unet(
                sample=sample[half_bs:],
                timestep=timestep,
                encoder_hidden_states=encoder_hidden_states[half_bs:],
                class_labels=class_labels,
                timestep_cond=timestep_cond,
                attention_mask=attention_mask,
                cross_attention_kwargs=cross_attention_kwargs,
                added_cond_kwargs=added_cond_kwargs,
                down_block_additional_residuals=down_block_additional_residuals_half2,
                mid_block_additional_residual=mid_block_additional_residual[half_bs:],
                encoder_attention_mask=encoder_attention_mask,
                return_dict=return_dict,
            )

            return (torch.cat([half1[0], half2[0]]),)

class AntiOOMAutoencoderKL(torch.nn.Module):
    def __init__(self, vae: AutoencoderKL):
        super(AntiOOMAutoencoderKL, self).__init__()
        self.vae = vae
        self.config = vae.config
        self.device = vae.device
        self.dtype = vae.dtype

    def modules(self):
        return self.vae.modules()

    def parameters(self, recurse: bool = True):
        return self.vae.parameters(recurse=recurse)

    def named_parameters(
        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
    ):
        return self.vae.named_parameters(
            prefix=prefix, recurse=recurse, remove_duplicate=remove_duplicate
        )
        
    def encode(
            self, 
            x: torch.FloatTensor, 
            return_dict: bool = True
        ):
        try:
            return self.vae.encode(
                x=x,
                return_dict=return_dict
            )
        except Exception as err:
            if "MUSA out of memory. " in str(err) or "CUDA out of memory. " in str(err):
                torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache() 
                gc.collect()
                # print("We need enable vae tiling against \"out of memory\"")
                self.vae.enable_tiling()
                result = self.vae.encode(
                    x=x,
                    return_dict=return_dict
                )
                
                return result
    
    def decode(
        self,
        z: torch.FloatTensor, 
        return_dict: bool = True
    ):
        try:
            return self.vae.decode(
                z=z,
                return_dict=return_dict
            )
        except Exception as err:
            if "MUSA out of memory. " in str(err) or "CUDA out of memory. " in str(err):
                torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache() 
                gc.collect()
                # print("We need enable vae tiling against \"out of memory\"")
                self.vae.enable_tiling()
                result = self.vae.decode(
                    z=z,
                    return_dict=return_dict
                )
                self.vae.disable_tiling()
                
                return result
modules/tools/anti_oom.py
MooreThreads/MobiMaliangSDK
import base64
from io import BytesIO
from PIL import Image

# given an image path, read it and return the base64 encoding
def read_img_as_base64(path):
    with open(path, "rb") as fr:
        encoded = base64.b64encode(fr.read()).decode()
        img = f"data:image/jpeg;base64,{encoded}"
    return img


# read the generation parameters from png info
def read_png_info(image):
    args = {}

    png_info = image.info
    params = png_info.pop("parameters")

    params = params.split("\nNegative prompt: ")
    args["prompt"] = params[0]

    params = params[1]
    pos = params.rfind("\n")
    args["negative_prompt"] = params[:pos]

    params = params[pos + 1 :].split(", ")

    args["steps"] = int(params[0].lstrip("Steps: "))
    args["cfg"] = float(params[2].lstrip("CFG scale: "))
    args["seed"] = int(params[3].lstrip("Seed: "))
    args["width"], args["height"] = params[4].lstrip("Size: ").split("x")
    args["width"] = int(args["width"])
    args["height"] = int(args["height"])

    return args

def image2base64(img):
    output_buffer = BytesIO()
    img.save(output_buffer, format='PNG')
    byte_data = output_buffer.getvalue()
    base64_str = base64.b64encode(byte_data)
    return base64_str

def base642image(base64_str):
    image = base64.b64decode(base64_str)
    image = BytesIO(image)
    image = Image.open(image)
    return image
modules/tools/image_utils.py
MooreThreads/MobiMaliangSDK
import logging
from logging.handlers import TimedRotatingFileHandler
import os
import sys
__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))

logdir = "logdirs"

class Logger():
    def __init__(self, logger_name='server_logger', logdir=logdir):
        
        self.logger = logging.getLogger(logger_name)
        self.logger.setLevel(logging.DEBUG)
        self.logdir = logdir
        os.makedirs(self.logdir,exist_ok=True)
        self.logfilename = os.path.join(self.logdir, 'server.log')
        
        
        self.consoleHandler = logging.StreamHandler()
        self.consoleHandler.setLevel(logging.WARNING)

        self.fileHandler = TimedRotatingFileHandler(filename=self.logfilename, when="D", backupCount=7, encoding="utf-8")
        self.fileHandler.setLevel(logging.INFO)
        
        self.formatter = logging.Formatter("%(asctime)s - %(filename)s ->\t%(funcName)s[line:%(lineno)d] - %(levelname)s: %(message)s")

        self.consoleHandler.setFormatter(self.formatter)
        self.fileHandler.setFormatter(self.formatter)

        self.logger.addHandler(self.consoleHandler)
        self.logger.addHandler(self.fileHandler)
        
    def getlog(self):
        
        self.logger.info("logger is loaded sucessfully!")
        return self.logger

server_logger = Logger(logdir=logdir).getlog()
modules/tools/loggers.py
MooreThreads/MobiMaliangSDK
import numpy as np


def get_enhancer(data, seed):
    np.random.seed(seed)
    quality = np.random.choice(
        data["quality"][0], size=3, replace=False, p=data["quality"][1]
    )

    composition = np.random.choice(
        data["composition"][0], size=1 if np.random.uniform(0, 1) > 0.5 else 0, replace=False, p=data["composition"][1]
    )

    color = np.random.choice(
        data["color"][0], size=1 if np.random.uniform(0, 1) > 0.5 else 0, replace=False, p=data["color"][1]
    )

    light = np.random.choice(
        data["light"][0], size=1, replace=False, p=data["light"][1]
    )

    content = np.random.choice(
        data["content"][0], size=3, replace=False, p=data["content"][1]
    )

    software = np.random.choice(
        data['software'][0], size=0, replace=False, p=data['software'][1]
    )

    photography = np.random.choice(
        data["photography"][0], size=0, replace=False, p=data["photography"][1]
    )

    result = np.concatenate(
        [quality, composition, color, light, content, software, photography]
    )
    np.random.shuffle(result)
    return ", ".join(list(result))


def get_artists(artists, seed):
    np.random.seed(seed)
    a = np.random.uniform(0, 1)

    # 0: 0.5; 1: 0.25; 2: 0.15; 3: 0.1
    if a < 0.5:
        num_artist = 0
        return ""
    elif a < 0.75:
        num_artist = 1
    elif a < 0.9:
        num_artist = 2
    else:
        num_artist = 3

    return ", ".join(np.random.choice(artists, num_artist, replace=False))


def get_neg_prompt():
    safety_neg_prompt = "nsfw, "
    neg_embeddings = "bad-picture-chill-75v, "
    quality_neg_prompt = "worst quality, low resolution, bad anatomy, malformed limbs, wrong fingers"
    return safety_neg_prompt, neg_embeddings, quality_neg_prompt

modules/tools/prompt_engineering.py
MooreThreads/MobiMaliangSDK
import re
from typing import List, Optional, Union

import torch

from diffusers import StableDiffusionPipeline
from diffusers.utils import logging

# ------------------------------------------------------------------------------

logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

re_attention = re.compile(
    r"""
\\\(|
\\\)|
\\\[|
\\]|
\\\\|
\\|
\(|
\[|
:([+-]?[.\d]+)\)|
\)|
]|
[^\\()\[\]:]+|
:
""",
    re.X,
)


def parse_prompt_attention(text):
    """
    Parses a string with attention tokens and returns a list of pairs: text and its associated weight.
    Accepted tokens are:
      (abc) - increases attention to abc by a multiplier of 1.1
      (abc:3.12) - increases attention to abc by a multiplier of 3.12
      [abc] - decreases attention to abc by a multiplier of 1.1
      \( - literal character '('
      \[ - literal character '['
      \) - literal character ')'
      \] - literal character ']'
      \\ - literal character '\'
      anything else - just text
    >>> parse_prompt_attention('normal text')
    [['normal text', 1.0]]
    >>> parse_prompt_attention('an (important) word')
    [['an ', 1.0], ['important', 1.1], [' word', 1.0]]
    >>> parse_prompt_attention('(unbalanced')
    [['unbalanced', 1.1]]
    >>> parse_prompt_attention('\(literal\]')
    [['(literal]', 1.0]]
    >>> parse_prompt_attention('(unnecessary)(parens)')
    [['unnecessaryparens', 1.1]]
    >>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).')
    [['a ', 1.0],
     ['house', 1.5730000000000004],
     [' ', 1.1],
     ['on', 1.0],
     [' a ', 1.1],
     ['hill', 0.55],
     [', sun, ', 1.1],
     ['sky', 1.4641000000000006],
     ['.', 1.1]]
    """

    res = []
    round_brackets = []
    square_brackets = []

    round_bracket_multiplier = 1.1
    square_bracket_multiplier = 1 / 1.1

    def multiply_range(start_position, multiplier):
        for p in range(start_position, len(res)):
            res[p][1] *= multiplier

    for m in re_attention.finditer(text):
        text = m.group(0)
        weight = m.group(1)

        if text.startswith("\\"):
            res.append([text[1:], 1.0])
        elif text == "(":
            round_brackets.append(len(res))
        elif text == "[":
            square_brackets.append(len(res))
        elif weight is not None and len(round_brackets) > 0:
            multiply_range(round_brackets.pop(), float(weight))
        elif text == ")" and len(round_brackets) > 0:
            multiply_range(round_brackets.pop(), round_bracket_multiplier)
        elif text == "]" and len(square_brackets) > 0:
            multiply_range(square_brackets.pop(), square_bracket_multiplier)
        else:
            res.append([text, 1.0])

    for pos in round_brackets:
        multiply_range(pos, round_bracket_multiplier)

    for pos in square_brackets:
        multiply_range(pos, square_bracket_multiplier)

    if len(res) == 0:
        res = [["", 1.0]]

    # merge runs of identical weights
    i = 0
    while i + 1 < len(res):
        if res[i][1] == res[i + 1][1]:
            res[i][0] += res[i + 1][0]
            res.pop(i + 1)
        else:
            i += 1

    return res


def get_prompts_with_weights(
    pipe: StableDiffusionPipeline, prompt: List[str], max_length: int
):
    r"""
    Tokenize a list of prompts and return its tokens with weights of each token.
    No padding, starting or ending token is included.
    """
    tokens = []
    weights = []
    truncated = False
    for text in prompt:
        texts_and_weights = parse_prompt_attention(text)
        text_token = []
        text_weight = []
        for word, weight in texts_and_weights:
            # tokenize and discard the starting and the ending token
            token = pipe.tokenizer(word).input_ids[1:-1]
            text_token += token
            # copy the weight by length of token
            text_weight += [weight] * len(token)
            # stop if the text is too long (longer than truncation limit)
            if len(text_token) > max_length:
                truncated = True
                break
        # truncate
        if len(text_token) > max_length:
            truncated = True
            text_token = text_token[:max_length]
            text_weight = text_weight[:max_length]
        tokens.append(text_token)
        weights.append(text_weight)
    if truncated:
        logger.warning(
            "Prompt was truncated. Try to shorten the prompt or increase max_embeddings_multiples"
        )
    return tokens, weights


def pad_tokens_and_weights(
    tokens, weights, max_length, bos, eos, no_boseos_middle=True, chunk_length=77
):
    r"""
    Pad the tokens (with starting and ending tokens) and weights (with 1.0) to max_length.
    """
    max_embeddings_multiples = (max_length - 2) // (chunk_length - 2)
    weights_length = (
        max_length if no_boseos_middle else max_embeddings_multiples * chunk_length
    )
    for i in range(len(tokens)):
        tokens[i] = [bos] + tokens[i] + [eos] * (max_length - 1 - len(tokens[i]))
        if no_boseos_middle:
            weights[i] = [1.0] + weights[i] + [1.0] * (max_length - 1 - len(weights[i]))
        else:
            w = []
            if len(weights[i]) == 0:
                w = [1.0] * weights_length
            else:
                for j in range(max_embeddings_multiples):
                    w.append(1.0)  # weight for starting token in this chunk
                    w += weights[i][
                        j
                        * (chunk_length - 2) : min(
                            len(weights[i]), (j + 1) * (chunk_length - 2)
                        )
                    ]
                    w.append(1.0)  # weight for ending token in this chunk
                w += [1.0] * (weights_length - len(w))
            weights[i] = w[:]

    return tokens, weights


def get_unweighted_text_embeddings(
    pipe: StableDiffusionPipeline,
    text_input: torch.Tensor,
    chunk_length: int,
    no_boseos_middle: Optional[bool] = True,
):
    """
    When the length of tokens is a multiple of the capacity of the text encoder,
    it should be split into chunks and sent to the text encoder individually.
    """
    max_embeddings_multiples = (text_input.shape[1] - 2) // (chunk_length - 2)
    if max_embeddings_multiples > 1:
        text_embeddings = []
        for i in range(max_embeddings_multiples):
            # extract the i-th chunk
            text_input_chunk = text_input[
                :, i * (chunk_length - 2) : (i + 1) * (chunk_length - 2) + 2
            ].clone()

            # cover the head and the tail by the starting and the ending tokens
            text_input_chunk[:, 0] = text_input[0, 0]
            text_input_chunk[:, -1] = text_input[0, -1]
            text_embedding = pipe.text_encoder(text_input_chunk)[0]

            if no_boseos_middle:
                if i == 0:
                    # discard the ending token
                    text_embedding = text_embedding[:, :-1]
                elif i == max_embeddings_multiples - 1:
                    # discard the starting token
                    text_embedding = text_embedding[:, 1:]
                else:
                    # discard both starting and ending tokens
                    text_embedding = text_embedding[:, 1:-1]

            text_embeddings.append(text_embedding)
        text_embeddings = torch.concat(text_embeddings, axis=1)
    else:
        text_embeddings = pipe.text_encoder(text_input)[0]
    return text_embeddings


def get_weighted_text_embeddings(
    pipe: StableDiffusionPipeline,
    prompt: Union[str, List[str]],
    uncond_prompt: Optional[Union[str, List[str]]] = None,
    max_embeddings_multiples: Optional[int] = 3,
    no_boseos_middle: Optional[bool] = False,
    skip_parsing: Optional[bool] = False,
    skip_weighting: Optional[bool] = False,
    **kwargs,
):
    r"""
    Prompts can be assigned with local weights using brackets. For example,
    prompt 'A (very beautiful) masterpiece' highlights the words 'very beautiful',
    and the embedding tokens corresponding to the words get multiplied by a constant, 1.1.
    Also, to regularize of the embedding, the weighted embedding would be scaled to preserve the original mean.
    Args:
        pipe (`StableDiffusionPipeline`):
            Pipe to provide access to the tokenizer and the text encoder.
        prompt (`str` or `List[str]`):
            The prompt or prompts to guide the image generation.
        uncond_prompt (`str` or `List[str]`):
            The unconditional prompt or prompts for guide the image generation. If unconditional prompt
            is provided, the embeddings of prompt and uncond_prompt are concatenated.
        max_embeddings_multiples (`int`, *optional*, defaults to `3`):
            The max multiple length of prompt embeddings compared to the max output length of text encoder.
        no_boseos_middle (`bool`, *optional*, defaults to `False`):
            If the length of text token is multiples of the capacity of text encoder, whether reserve the starting and
            ending token in each of the chunk in the middle.
        skip_parsing (`bool`, *optional*, defaults to `False`):
            Skip the parsing of brackets.
        skip_weighting (`bool`, *optional*, defaults to `False`):
            Skip the weighting. When the parsing is skipped, it is forced True.
    """
    max_length = (pipe.tokenizer.model_max_length - 2) * max_embeddings_multiples + 2
    if isinstance(prompt, str):
        prompt = [prompt]

    if not skip_parsing:
        prompt_tokens, prompt_weights = get_prompts_with_weights(
            pipe, prompt, max_length - 2
        )
        if uncond_prompt is not None:
            if isinstance(uncond_prompt, str):
                uncond_prompt = [uncond_prompt]
            uncond_tokens, uncond_weights = get_prompts_with_weights(
                pipe, uncond_prompt, max_length - 2
            )
    else:
        prompt_tokens = [
            token[1:-1]
            for token in pipe.tokenizer(
                prompt, max_length=max_length, truncation=True
            ).input_ids
        ]
        prompt_weights = [[1.0] * len(token) for token in prompt_tokens]
        if uncond_prompt is not None:
            if isinstance(uncond_prompt, str):
                uncond_prompt = [uncond_prompt]
            uncond_tokens = [
                token[1:-1]
                for token in pipe.tokenizer(
                    uncond_prompt, max_length=max_length, truncation=True
                ).input_ids
            ]
            uncond_weights = [[1.0] * len(token) for token in uncond_tokens]

    # round up the longest length of tokens to a multiple of (model_max_length - 2)
    max_length = max([len(token) for token in prompt_tokens])
    if uncond_prompt is not None:
        max_length = max(max_length, max([len(token) for token in uncond_tokens]))

    max_embeddings_multiples = min(
        max_embeddings_multiples,
        (max_length - 1) // (pipe.tokenizer.model_max_length - 2) + 1,
    )
    max_embeddings_multiples = max(1, max_embeddings_multiples)
    max_length = (pipe.tokenizer.model_max_length - 2) * max_embeddings_multiples + 2

    # pad the length of tokens and weights
    bos = pipe.tokenizer.bos_token_id
    eos = pipe.tokenizer.eos_token_id
    prompt_tokens, prompt_weights = pad_tokens_and_weights(
        prompt_tokens,
        prompt_weights,
        max_length,
        bos,
        eos,
        no_boseos_middle=no_boseos_middle,
        chunk_length=pipe.tokenizer.model_max_length,
    )
    prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.long, device=pipe.device)
    if uncond_prompt is not None:
        uncond_tokens, uncond_weights = pad_tokens_and_weights(
            uncond_tokens,
            uncond_weights,
            max_length,
            bos,
            eos,
            no_boseos_middle=no_boseos_middle,
            chunk_length=pipe.tokenizer.model_max_length,
        )
        uncond_tokens = torch.tensor(
            uncond_tokens, dtype=torch.long, device=pipe.device
        )

    # get the embeddings
    text_embeddings = get_unweighted_text_embeddings(
        pipe,
        prompt_tokens,
        pipe.tokenizer.model_max_length,
        no_boseos_middle=no_boseos_middle,
    )
    prompt_weights = torch.tensor(
        prompt_weights, dtype=text_embeddings.dtype, device=pipe.device
    )
    if uncond_prompt is not None:
        uncond_embeddings = get_unweighted_text_embeddings(
            pipe,
            uncond_tokens,
            pipe.tokenizer.model_max_length,
            no_boseos_middle=no_boseos_middle,
        )
        uncond_weights = torch.tensor(
            uncond_weights, dtype=uncond_embeddings.dtype, device=pipe.device
        )

    # assign weights to the prompts and normalize in the sense of mean
    # TODO: should we normalize by chunk or in a whole (current implementation)?
    if (not skip_parsing) and (not skip_weighting):
        previous_mean = (
            text_embeddings.float().mean(axis=[-2, -1]).to(text_embeddings.dtype)
        )
        text_embeddings *= prompt_weights.unsqueeze(-1)
        current_mean = (
            text_embeddings.float().mean(axis=[-2, -1]).to(text_embeddings.dtype)
        )
        text_embeddings *= (previous_mean / current_mean).unsqueeze(-1).unsqueeze(-1)
        if uncond_prompt is not None:
            previous_mean = (
                uncond_embeddings.float()
                .mean(axis=[-2, -1])
                .to(uncond_embeddings.dtype)
            )
            uncond_embeddings *= uncond_weights.unsqueeze(-1)
            current_mean = (
                uncond_embeddings.float()
                .mean(axis=[-2, -1])
                .to(uncond_embeddings.dtype)
            )
            uncond_embeddings *= (
                (previous_mean / current_mean).unsqueeze(-1).unsqueeze(-1)
            )

    if uncond_prompt is not None:
        return text_embeddings, uncond_embeddings
    return text_embeddings, None


# get names and weights of loras from a prompt in the format of <lora:name:weight>
def get_lora_from_prompt(prompt):
    pattern = "<lora:\S*:\S*>"
    final_prompt = ""
    result = []
    res = re.search(pattern, prompt)
    while res is not None:
        lora = res.group()
        idx = prompt.find(lora)
        final_prompt += prompt[:idx]
        prompt = prompt[idx + len(lora) :]
        lora = lora.split(":")
        try:
            lora_name = lora[1]
            lora_weight = float(lora[2][:-1])
        except:
            pass
        else:
            result.append([lora_name, lora_weight])

        res = re.search(pattern, prompt)

    final_prompt += prompt

    return result, final_prompt


# whether a prompt contains chinese or not
def contain_zh(s):
    for w in s:
        if "\u4e00" <= w <= "\u9fff":
            return True

    return False

modules/tools/prompt_parser.py
MooreThreads/MobiMaliangSDK
import torch
import numpy as np
from diffusers import DDIMScheduler
from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput
from diffusers.utils import randn_tensor
from typing import List, Optional, Tuple, Union


class BiDDIMScheduler(DDIMScheduler):
    def set_timesteps(
        self,
        num_inference_steps: int,
        device: Union[str, torch.device] = None,
        bi_sampler_start_ratio: float = 0.2,
    ):
        """
        Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.

        Args:
            num_inference_steps (`int`):
                the number of diffusion steps used when generating samples with a pre-trained model.
        """

        assert (
            0 <= bi_sampler_start_ratio < 0.5
        ), "please set bi_sampler_start_ratio less than 0.5 and larger than 0"
        if num_inference_steps > self.config.num_train_timesteps:
            raise ValueError(
                f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
                f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
                f" maximal {self.config.num_train_timesteps} timesteps."
            )

        self.num_inference_steps = num_inference_steps
        self.bi_sampler_start_ratio = bi_sampler_start_ratio

        # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
        if self.config.timestep_spacing == "linspace":
            timesteps = (
                np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
                .round()[::-1]
                .copy()
                .astype(np.int64)
            )
        elif self.config.timestep_spacing == "leading":
            step_ratio = self.config.num_train_timesteps // self.num_inference_steps
            # creates integer timesteps by multiplying by ratio
            # casting to int to avoid issues when num_inference_step is power of 3
            timesteps1 = (
                (np.arange(0, num_inference_steps) * step_ratio)
                .round()[::-1]
                .copy()
                .astype(np.int64)
            )
            timesteps2 = (
                (np.arange(0, num_inference_steps // 2) * step_ratio * 2)
                .round()[::-1]
                .copy()
                .astype(np.int64)
            )
            timesteps = []
            for i in range(0, round(len(timesteps1) * bi_sampler_start_ratio)):
                timesteps.append(timesteps1[i])
            for i in range(
                round(len(timesteps2) * bi_sampler_start_ratio),
                round(len(timesteps2) * (1 - bi_sampler_start_ratio)),
            ):
                timesteps.append(timesteps2[i])
            for i in range(
                round(len(timesteps1) * (1 - bi_sampler_start_ratio)), len(timesteps1)
            ):
                timesteps.append(timesteps1[i])
            timesteps = np.array(timesteps).copy().astype(np.int64)
            timesteps += self.config.steps_offset
        elif self.config.timestep_spacing == "trailing":
            step_ratio = self.config.num_train_timesteps / self.num_inference_steps
            # creates integer timesteps by multiplying by ratio
            # casting to int to avoid issues when num_inference_step is power of 3
            timesteps = np.round(
                np.arange(self.config.num_train_timesteps, 0, -step_ratio)
            ).astype(np.int64)
            timesteps -= 1
        else:
            raise ValueError(
                f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'."
            )

        self.timesteps = torch.from_numpy(timesteps).to(device)

    def step(
        self,
        model_output: torch.FloatTensor,
        timestep: int,
        sample: torch.FloatTensor,
        eta: float = 0.0,
        use_clipped_model_output: bool = False,
        generator=None,
        variance_noise: Optional[torch.FloatTensor] = None,
        return_dict: bool = True,
    ) -> Union[DDIMSchedulerOutput, Tuple]:
        """
        Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
        process from the learned model outputs (most often the predicted noise).

        Args:
            model_output (`torch.FloatTensor`): direct output from learned diffusion model.
            timestep (`int`): current discrete timestep in the diffusion chain.
            sample (`torch.FloatTensor`):
                current instance of sample being created by diffusion process.
            eta (`float`): weight of noise for added noise in diffusion step.
            use_clipped_model_output (`bool`): if `True`, compute "corrected" `model_output` from the clipped
                predicted original sample. Necessary because predicted original sample is clipped to [-1, 1] when
                `self.config.clip_sample` is `True`. If no clipping has happened, "corrected" `model_output` would
                coincide with the one provided as input and `use_clipped_model_output` will have not effect.
            generator: random number generator.
            variance_noise (`torch.FloatTensor`): instead of generating noise for the variance using `generator`, we
                can directly provide the noise for the variance itself. This is useful for methods such as
                CycleDiffusion. (https://arxiv.org/abs/2210.05559)
            return_dict (`bool`): option for returning tuple rather than DDIMSchedulerOutput class

        Returns:
            [`~schedulers.scheduling_utils.DDIMSchedulerOutput`] or `tuple`:
            [`~schedulers.scheduling_utils.DDIMSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When
            returning a tuple, the first element is the sample tensor.

        """
        if self.num_inference_steps is None:
            raise ValueError(
                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
            )

        # 1. get previous step value (=t-1)
        if (
            self.bi_sampler_start_ratio
            <= (timestep.to("cpu") - self.config.steps_offset)
            / self.config.num_train_timesteps
            < 1 - self.bi_sampler_start_ratio
        ):
            prev_timestep = (
                timestep
                - 2 * self.config.num_train_timesteps // self.num_inference_steps
            )
        else:
            prev_timestep = (
                timestep - self.config.num_train_timesteps // self.num_inference_steps
            )
        # 2. compute alphas, betas
        alpha_prod_t = self.alphas_cumprod[timestep]
        alpha_prod_t_prev = (
            self.alphas_cumprod[prev_timestep]
            if prev_timestep >= 0
            else self.final_alpha_cumprod
        )

        beta_prod_t = 1 - alpha_prod_t

        # 3. compute predicted original sample from predicted noise also called
        # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        if self.config.prediction_type == "epsilon":
            pred_original_sample = (
                sample - beta_prod_t ** (0.5) * model_output
            ) / alpha_prod_t ** (0.5)
            pred_epsilon = model_output
        elif self.config.prediction_type == "sample":
            pred_original_sample = model_output
            pred_epsilon = (
                sample - alpha_prod_t ** (0.5) * pred_original_sample
            ) / beta_prod_t ** (0.5)
        elif self.config.prediction_type == "v_prediction":
            pred_original_sample = (alpha_prod_t ** 0.5) * sample - (
                beta_prod_t ** 0.5
            ) * model_output
            pred_epsilon = (alpha_prod_t ** 0.5) * model_output + (
                beta_prod_t ** 0.5
            ) * sample
        else:
            raise ValueError(
                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
                " `v_prediction`"
            )

        # 4. Clip or threshold "predicted x_0"
        if self.config.thresholding:
            pred_original_sample = self._threshold_sample(pred_original_sample)
        elif self.config.clip_sample:
            pred_original_sample = pred_original_sample.clamp(
                -self.config.clip_sample_range, self.config.clip_sample_range
            )

        # 5. compute variance: "sigma_t(η)" -> see formula (16)
        # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
        variance = self._get_variance(timestep, prev_timestep)
        std_dev_t = eta * variance ** (0.5)

        if use_clipped_model_output:
            # the pred_epsilon is always re-derived from the clipped x_0 in Glide
            pred_epsilon = (
                sample - alpha_prod_t ** (0.5) * pred_original_sample
            ) / beta_prod_t ** (0.5)

        # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t ** 2) ** (
            0.5
        ) * pred_epsilon

        # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        prev_sample = (
            alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
        )

        if eta > 0:
            if variance_noise is not None and generator is not None:
                raise ValueError(
                    "Cannot pass both generator and variance_noise. Please make sure that either `generator` or"
                    " `variance_noise` stays `None`."
                )

            if variance_noise is None:
                variance_noise = randn_tensor(
                    model_output.shape,
                    generator=generator,
                    device=model_output.device,
                    dtype=model_output.dtype,
                )
            variance = std_dev_t * variance_noise

            prev_sample = prev_sample + variance

        if not return_dict:
            return (prev_sample,)

        return DDIMSchedulerOutput(
            prev_sample=prev_sample, pred_original_sample=pred_original_sample
        )

modules/tools/samplers.py
MooreThreads/MobiMaliangSDK
import numpy as np
from PIL import Image
from tqdm import tqdm

import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from wrappers import RealESRGANWrapper


class TileRealESRGAN(RealESRGANWrapper):
    def __init__(self, device):
        super(TileRealESRGAN, self).__init__(device=device)

    def get_views(self, width, height, window_size=720):
        stride = window_size - 2
        window_height_size = min(height, window_size)
        window_width_size = min(width, window_size)

        if height % stride == 0:
            num_blocks_height = int(max(1, height // stride))
        else:
            num_blocks_height = int(height // stride + 1)
        if width % stride == 0:
            num_blocks_width = int(max(1, width // stride))
        else:
            num_blocks_width = int(width // stride + 1)
        total_num_blocks = int(num_blocks_height * num_blocks_width)
        views = []

        for i in range(total_num_blocks):
            h_start = int((i // num_blocks_width) * stride)
            h_end = h_start + window_height_size
            if h_end > height:
                h_start = int(height - stride)
                h_end = int(height)
            w_start = int((i % num_blocks_width) * stride)
            w_end = w_start + window_width_size
            if w_end > width:
                w_start = int(width - stride)
                w_end = int(width)
            views.append((h_start, h_end, w_start, w_end))
        return views

    def __call__(self, image):

        self.start()

        sr_ratio = 4
        width, height = image.size
        width = int(width // 2 * 2)
        height = int(height // 2 * 2)
        image = image.resize((width, height))
        view_grids = self.get_views(width, height)
        np_value = np.zeros((height * sr_ratio, width * sr_ratio, 3))
        np_count = np.zeros_like(np_value)

        for h_start, h_end, w_start, w_end in tqdm(view_grids):
            image_crop = image.crop((w_start, h_start, w_end, h_end))
            sr_image_crop = self.super_resolution(image_crop)
            np_sr_image_crop = np.array(sr_image_crop)

            np_value[
                h_start * sr_ratio : h_end * sr_ratio,
                w_start * sr_ratio : w_end * sr_ratio,
            ] += np_sr_image_crop
            np_count[
                h_start * sr_ratio : h_end * sr_ratio,
                w_start * sr_ratio : w_end * sr_ratio,
            ] += 1

        self.stop()
        np_sr_image = np.where(np_count > 0, np_value / np_count, np_value).astype(
            np.uint8
        )
        sr_image = Image.fromarray(np_sr_image)

        return sr_image

modules/tools/tile_sr.py
MooreThreads/MobiMaliangSDK
from .controlnet_annotator import (
    CannyDetectorWrapper,
    ContentShuffleDetectorWrapper,
    InpaintingConditionalImageWrapper,
    ResizeForConditionalImageWrapper,
    ZoomConditionalImageWrapper,
    OpenposeDetectorWrapper,
    LineartDetectorWrapper,
    NormalBaeDetectorWrapper,
    HEDdetectorWrapper,
    PidiNetDetectorWrapper,
    MLSDdetectorWrapper,
)
from .deepbooru import DeepDanbooruWrapper
from .realesrgan import RealESRGANWrapper
from .translator import Translator

modules/wrappers/__init__.py
MooreThreads/MobiMaliangSDK
import torch
import gc
import numpy as np
from PIL import Image
from controlnet_aux import (
    CannyDetector,
    ContentShuffleDetector,
    OpenposeDetector,
    LineartDetector,
    NormalBaeDetector,
    HEDdetector,
    PidiNetDetector,
    MLSDdetector,
)
from transformers import pipeline

import os
import sys
__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from shared import MTGPU_DETECTION, device
    
class CannyDetectorWrapper:
    def __init__(self):
        self.processor = CannyDetector()

    def __call__(self, input_image, low_threshold=100, high_threshold=200):
        return self.processor(input_image, low_threshold, high_threshold)


class ContentShuffleDetectorWrapper:
    def __init__(self):
        self.processor = ContentShuffleDetector()

    def __call__(
        self,
        input_image,
        detect_resolution=512,
        image_resolution=512,
        output_type="pil",
    ):
        return self.processor(
            input_image=input_image,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            output_type=output_type,
        )


class ResizeForConditionalImageWrapper:
    def __init__(self):
        pass

    def __call__(self, input_image, down_sampling_ratio=8.0):
        input_image = input_image.convert("RGB")
        W, H = input_image.size
        w, h = int(W // down_sampling_ratio), int(H // down_sampling_ratio)
        img = input_image.resize((w, h), resample=Image.LANCZOS)
        H = int(round(H / 8.0)) * 8
        W = int(round(W / 8.0)) * 8
        img = img.resize((W, H), resample=Image.LANCZOS)
        return img


class ZoomConditionalImageWrapper:
    def __init__(self):
        pass

    def __call__(self, input_image, up_sampling_ratio=2.0):
        input_image = input_image.convert("RGB")
        w, h = input_image.size
        W, H = int(w * up_sampling_ratio), int(h * up_sampling_ratio)
        H = int(round(H / 8.0)) * 8
        W = int(round(W / 8.0)) * 8
        img = input_image.resize((W, H), resample=Image.LANCZOS)
        return img


class InpaintingConditionalImageWrapper:
    def __init__(self):
        pass

    def __call__(self, input_image, mask_image):

        image = np.array(input_image.convert("RGB")).astype(np.float32) / 255.0
        mask_image = np.array(mask_image.convert("L")).astype(np.float32) / 255.0

        assert (
            image.shape[0:1] == mask_image.shape[0:1]
        ), "image and mask_image must have the same image size"
        image[mask_image > 0.5] = -1.0  # set as masked pixel
        image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
        image = torch.from_numpy(image)
        return image


class ContorlnetAnnotatorBasicWrapper:
    def __init__(self, device=device):
        self.device = device
        self.processor = None

    def start(self):
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()
        gc.collect()
        assert self.processor is not None, "pleaase load processor first"
        self.processor.to(self.device)

    def stop(self):
        self.processor.to("cpu")
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()
        gc.collect()


class OpenposeDetectorWrapper(ContorlnetAnnotatorBasicWrapper):
    def __init__(self, device=device):
        super(OpenposeDetectorWrapper, self).__init__(device=device)

    def load_models(self, models_dir):
        self.processor = OpenposeDetector.from_pretrained(models_dir).to("cpu")

    def __call__(self, input_image, include_hand=True, include_face=True):

        self.start()
        procesed_image = self.processor(
            input_image=input_image,
            include_hand=include_hand,
            include_face=include_face,
        )
        self.stop()

        return procesed_image


class LineartDetectorWrapper(ContorlnetAnnotatorBasicWrapper):
    def __init__(self, device=device):
        super(LineartDetectorWrapper, self).__init__(device=device)

    def load_models(self, models_dir):
        self.processor = LineartDetector.from_pretrained(models_dir).to("cpu")

    def __call__(
        self,
        input_image,
        coarse=False,
        detect_resolution=512,
        image_resolution=512,
        output_type="pil",
    ):

        self.start()
        procesed_image = self.processor(
            input_image=input_image,
            coarse=coarse,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            output_type=output_type,
        )
        self.stop()

        return procesed_image


class NormalBaeDetectorWrapper(ContorlnetAnnotatorBasicWrapper):
    def __init__(self, device=device):
        super(NormalBaeDetectorWrapper, self).__init__(device=device)

    def load_models(self, models_dir):
        self.processor = NormalBaeDetector.from_pretrained(models_dir).to("cpu")

    def __call__(
        self,
        input_image,
        detect_resolution=512,
        image_resolution=512,
        output_type="pil",
    ):

        self.start()
        procesed_image = self.processor(
            input_image=input_image,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            output_type=output_type,
        )
        self.stop()

        return procesed_image


class HEDdetectorWrapper(ContorlnetAnnotatorBasicWrapper):
    def __init__(self, device=device):
        super(HEDdetectorWrapper, self).__init__(device=device)

    def load_models(self, models_dir):
        self.processor = HEDdetector.from_pretrained(models_dir).to("cpu")

    def __call__(
        self,
        input_image,
        detect_resolution=512,
        image_resolution=512,
        safe=False,
        output_type="pil",
        scribble=False,
    ):

        self.start()
        procesed_image = self.processor(
            input_image=input_image,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            safe=safe,
            output_type=output_type,
            scribble=scribble,
        )
        self.stop()

        return procesed_image


class PidiNetDetectorWrapper(ContorlnetAnnotatorBasicWrapper):
    def __init__(self, device=device):
        super(PidiNetDetectorWrapper, self).__init__(device=device)

    def load_models(self, models_dir):
        self.processor = PidiNetDetector.from_pretrained(models_dir).to("cpu")

    def __call__(
        self,
        input_image,
        detect_resolution=512,
        image_resolution=512,
        safe=False,
        output_type="pil",
        scribble=False,
        apply_filter=False,
    ):

        self.start()
        procesed_image = self.processor(
            input_image=input_image,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            safe=safe,
            output_type=output_type,
            scribble=scribble,
            apply_filter=apply_filter,
        )
        self.stop()

        return procesed_image


class MLSDdetectorWrapper(ContorlnetAnnotatorBasicWrapper):
    def __init__(self, device=device):
        super(MLSDdetectorWrapper, self).__init__(device=device)

    def load_models(self, models_dir):
        self.processor = MLSDdetector.from_pretrained(models_dir).to("cpu")

    def __call__(
        self,
        input_image,
        thr_v=0.1,
        thr_d=0.1,
        detect_resolution=512,
        image_resolution=512,
        output_type="pil",
    ):

        self.start()
        procesed_image = self.processor(
            input_image=input_image,
            thr_v=thr_v,
            thr_d=thr_d,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
            output_type=output_type,
        )
        self.stop()

        return procesed_image

modules/wrappers/controlnet_annotator.py
MooreThreads/MobiMaliangSDK
import torch
import torch.nn as nn
import gc
import torch.nn.functional as F
import numpy as np
import re

import os
import sys
__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.shared import MTGPU_DETECTION

unet_needs_upcast = False
re_special = re.compile(r"([\\()])")


class DeepDanbooruModel(nn.Module):
    def __init__(self):
        super(DeepDanbooruModel, self).__init__()

        self.tags = []

        self.n_Conv_0 = nn.Conv2d(
            kernel_size=(7, 7), in_channels=3, out_channels=64, stride=(2, 2)
        )
        self.n_MaxPool_0 = nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2))
        self.n_Conv_1 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256)
        self.n_Conv_2 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=64)
        self.n_Conv_3 = nn.Conv2d(kernel_size=(3, 3), in_channels=64, out_channels=64)
        self.n_Conv_4 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256)
        self.n_Conv_5 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=64)
        self.n_Conv_6 = nn.Conv2d(kernel_size=(3, 3), in_channels=64, out_channels=64)
        self.n_Conv_7 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256)
        self.n_Conv_8 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=64)
        self.n_Conv_9 = nn.Conv2d(kernel_size=(3, 3), in_channels=64, out_channels=64)
        self.n_Conv_10 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256)
        self.n_Conv_11 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=512, stride=(2, 2)
        )
        self.n_Conv_12 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=128
        )
        self.n_Conv_13 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=128, out_channels=128, stride=(2, 2)
        )
        self.n_Conv_14 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=128, out_channels=512
        )
        self.n_Conv_15 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=128
        )
        self.n_Conv_16 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=128, out_channels=128
        )
        self.n_Conv_17 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=128, out_channels=512
        )
        self.n_Conv_18 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=128
        )
        self.n_Conv_19 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=128, out_channels=128
        )
        self.n_Conv_20 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=128, out_channels=512
        )
        self.n_Conv_21 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=128
        )
        self.n_Conv_22 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=128, out_channels=128
        )
        self.n_Conv_23 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=128, out_channels=512
        )
        self.n_Conv_24 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=128
        )
        self.n_Conv_25 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=128, out_channels=128
        )
        self.n_Conv_26 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=128, out_channels=512
        )
        self.n_Conv_27 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=128
        )
        self.n_Conv_28 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=128, out_channels=128
        )
        self.n_Conv_29 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=128, out_channels=512
        )
        self.n_Conv_30 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=128
        )
        self.n_Conv_31 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=128, out_channels=128
        )
        self.n_Conv_32 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=128, out_channels=512
        )
        self.n_Conv_33 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=128
        )
        self.n_Conv_34 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=128, out_channels=128
        )
        self.n_Conv_35 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=128, out_channels=512
        )
        self.n_Conv_36 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=1024, stride=(2, 2)
        )
        self.n_Conv_37 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=256
        )
        self.n_Conv_38 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256, stride=(2, 2)
        )
        self.n_Conv_39 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_40 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_41 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_42 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_43 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_44 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_45 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_46 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_47 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_48 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_49 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_50 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_51 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_52 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_53 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_54 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_55 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_56 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_57 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_58 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_59 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_60 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_61 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_62 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_63 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_64 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_65 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_66 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_67 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_68 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_69 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_70 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_71 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_72 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_73 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_74 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_75 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_76 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_77 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_78 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_79 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_80 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_81 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_82 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_83 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_84 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_85 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_86 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_87 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_88 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_89 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_90 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_91 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_92 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_93 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_94 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_95 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_96 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_97 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_98 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256, stride=(2, 2)
        )
        self.n_Conv_99 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_100 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=1024, stride=(2, 2)
        )
        self.n_Conv_101 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_102 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_103 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_104 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_105 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_106 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_107 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_108 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_109 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_110 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_111 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_112 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_113 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_114 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_115 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_116 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_117 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_118 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_119 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_120 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_121 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_122 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_123 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_124 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_125 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_126 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_127 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_128 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_129 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_130 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_131 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_132 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_133 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_134 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_135 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_136 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_137 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_138 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_139 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_140 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_141 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_142 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_143 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_144 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_145 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_146 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_147 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_148 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_149 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_150 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_151 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_152 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_153 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_154 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_155 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=256
        )
        self.n_Conv_156 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=256, out_channels=256
        )
        self.n_Conv_157 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=256, out_channels=1024
        )
        self.n_Conv_158 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=2048, stride=(2, 2)
        )
        self.n_Conv_159 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=512
        )
        self.n_Conv_160 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=512, out_channels=512, stride=(2, 2)
        )
        self.n_Conv_161 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=2048
        )
        self.n_Conv_162 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=2048, out_channels=512
        )
        self.n_Conv_163 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=512, out_channels=512
        )
        self.n_Conv_164 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=2048
        )
        self.n_Conv_165 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=2048, out_channels=512
        )
        self.n_Conv_166 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=512, out_channels=512
        )
        self.n_Conv_167 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=512, out_channels=2048
        )
        self.n_Conv_168 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=2048, out_channels=4096, stride=(2, 2)
        )
        self.n_Conv_169 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=2048, out_channels=1024
        )
        self.n_Conv_170 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=1024, out_channels=1024, stride=(2, 2)
        )
        self.n_Conv_171 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=4096
        )
        self.n_Conv_172 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=4096, out_channels=1024
        )
        self.n_Conv_173 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=1024, out_channels=1024
        )
        self.n_Conv_174 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=4096
        )
        self.n_Conv_175 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=4096, out_channels=1024
        )
        self.n_Conv_176 = nn.Conv2d(
            kernel_size=(3, 3), in_channels=1024, out_channels=1024
        )
        self.n_Conv_177 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=1024, out_channels=4096
        )
        self.n_Conv_178 = nn.Conv2d(
            kernel_size=(1, 1), in_channels=4096, out_channels=9176, bias=False
        )

    def forward(self, *inputs):
        t_358, = inputs
        t_359 = t_358.permute(*[0, 3, 1, 2])
        t_359_padded = F.pad(t_359, [2, 3, 2, 3], value=0)
        t_360 = self.n_Conv_0(
            t_359_padded.to(self.n_Conv_0.bias.dtype)
            if unet_needs_upcast
            else t_359_padded
        )
        t_361 = F.relu(t_360)
        t_361 = F.pad(t_361, [0, 1, 0, 1], value=float("-inf"))
        t_362 = self.n_MaxPool_0(t_361)
        t_363 = self.n_Conv_1(t_362)
        t_364 = self.n_Conv_2(t_362)
        t_365 = F.relu(t_364)
        t_365_padded = F.pad(t_365, [1, 1, 1, 1], value=0)
        t_366 = self.n_Conv_3(t_365_padded)
        t_367 = F.relu(t_366)
        t_368 = self.n_Conv_4(t_367)
        t_369 = torch.add(t_368, t_363)
        t_370 = F.relu(t_369)
        t_371 = self.n_Conv_5(t_370)
        t_372 = F.relu(t_371)
        t_372_padded = F.pad(t_372, [1, 1, 1, 1], value=0)
        t_373 = self.n_Conv_6(t_372_padded)
        t_374 = F.relu(t_373)
        t_375 = self.n_Conv_7(t_374)
        t_376 = torch.add(t_375, t_370)
        t_377 = F.relu(t_376)
        t_378 = self.n_Conv_8(t_377)
        t_379 = F.relu(t_378)
        t_379_padded = F.pad(t_379, [1, 1, 1, 1], value=0)
        t_380 = self.n_Conv_9(t_379_padded)
        t_381 = F.relu(t_380)
        t_382 = self.n_Conv_10(t_381)
        t_383 = torch.add(t_382, t_377)
        t_384 = F.relu(t_383)
        t_385 = self.n_Conv_11(t_384)
        t_386 = self.n_Conv_12(t_384)
        t_387 = F.relu(t_386)
        t_387_padded = F.pad(t_387, [0, 1, 0, 1], value=0)
        t_388 = self.n_Conv_13(t_387_padded)
        t_389 = F.relu(t_388)
        t_390 = self.n_Conv_14(t_389)
        t_391 = torch.add(t_390, t_385)
        t_392 = F.relu(t_391)
        t_393 = self.n_Conv_15(t_392)
        t_394 = F.relu(t_393)
        t_394_padded = F.pad(t_394, [1, 1, 1, 1], value=0)
        t_395 = self.n_Conv_16(t_394_padded)
        t_396 = F.relu(t_395)
        t_397 = self.n_Conv_17(t_396)
        t_398 = torch.add(t_397, t_392)
        t_399 = F.relu(t_398)
        t_400 = self.n_Conv_18(t_399)
        t_401 = F.relu(t_400)
        t_401_padded = F.pad(t_401, [1, 1, 1, 1], value=0)
        t_402 = self.n_Conv_19(t_401_padded)
        t_403 = F.relu(t_402)
        t_404 = self.n_Conv_20(t_403)
        t_405 = torch.add(t_404, t_399)
        t_406 = F.relu(t_405)
        t_407 = self.n_Conv_21(t_406)
        t_408 = F.relu(t_407)
        t_408_padded = F.pad(t_408, [1, 1, 1, 1], value=0)
        t_409 = self.n_Conv_22(t_408_padded)
        t_410 = F.relu(t_409)
        t_411 = self.n_Conv_23(t_410)
        t_412 = torch.add(t_411, t_406)
        t_413 = F.relu(t_412)
        t_414 = self.n_Conv_24(t_413)
        t_415 = F.relu(t_414)
        t_415_padded = F.pad(t_415, [1, 1, 1, 1], value=0)
        t_416 = self.n_Conv_25(t_415_padded)
        t_417 = F.relu(t_416)
        t_418 = self.n_Conv_26(t_417)
        t_419 = torch.add(t_418, t_413)
        t_420 = F.relu(t_419)
        t_421 = self.n_Conv_27(t_420)
        t_422 = F.relu(t_421)
        t_422_padded = F.pad(t_422, [1, 1, 1, 1], value=0)
        t_423 = self.n_Conv_28(t_422_padded)
        t_424 = F.relu(t_423)
        t_425 = self.n_Conv_29(t_424)
        t_426 = torch.add(t_425, t_420)
        t_427 = F.relu(t_426)
        t_428 = self.n_Conv_30(t_427)
        t_429 = F.relu(t_428)
        t_429_padded = F.pad(t_429, [1, 1, 1, 1], value=0)
        t_430 = self.n_Conv_31(t_429_padded)
        t_431 = F.relu(t_430)
        t_432 = self.n_Conv_32(t_431)
        t_433 = torch.add(t_432, t_427)
        t_434 = F.relu(t_433)
        t_435 = self.n_Conv_33(t_434)
        t_436 = F.relu(t_435)
        t_436_padded = F.pad(t_436, [1, 1, 1, 1], value=0)
        t_437 = self.n_Conv_34(t_436_padded)
        t_438 = F.relu(t_437)
        t_439 = self.n_Conv_35(t_438)
        t_440 = torch.add(t_439, t_434)
        t_441 = F.relu(t_440)
        t_442 = self.n_Conv_36(t_441)
        t_443 = self.n_Conv_37(t_441)
        t_444 = F.relu(t_443)
        t_444_padded = F.pad(t_444, [0, 1, 0, 1], value=0)
        t_445 = self.n_Conv_38(t_444_padded)
        t_446 = F.relu(t_445)
        t_447 = self.n_Conv_39(t_446)
        t_448 = torch.add(t_447, t_442)
        t_449 = F.relu(t_448)
        t_450 = self.n_Conv_40(t_449)
        t_451 = F.relu(t_450)
        t_451_padded = F.pad(t_451, [1, 1, 1, 1], value=0)
        t_452 = self.n_Conv_41(t_451_padded)
        t_453 = F.relu(t_452)
        t_454 = self.n_Conv_42(t_453)
        t_455 = torch.add(t_454, t_449)
        t_456 = F.relu(t_455)
        t_457 = self.n_Conv_43(t_456)
        t_458 = F.relu(t_457)
        t_458_padded = F.pad(t_458, [1, 1, 1, 1], value=0)
        t_459 = self.n_Conv_44(t_458_padded)
        t_460 = F.relu(t_459)
        t_461 = self.n_Conv_45(t_460)
        t_462 = torch.add(t_461, t_456)
        t_463 = F.relu(t_462)
        t_464 = self.n_Conv_46(t_463)
        t_465 = F.relu(t_464)
        t_465_padded = F.pad(t_465, [1, 1, 1, 1], value=0)
        t_466 = self.n_Conv_47(t_465_padded)
        t_467 = F.relu(t_466)
        t_468 = self.n_Conv_48(t_467)
        t_469 = torch.add(t_468, t_463)
        t_470 = F.relu(t_469)
        t_471 = self.n_Conv_49(t_470)
        t_472 = F.relu(t_471)
        t_472_padded = F.pad(t_472, [1, 1, 1, 1], value=0)
        t_473 = self.n_Conv_50(t_472_padded)
        t_474 = F.relu(t_473)
        t_475 = self.n_Conv_51(t_474)
        t_476 = torch.add(t_475, t_470)
        t_477 = F.relu(t_476)
        t_478 = self.n_Conv_52(t_477)
        t_479 = F.relu(t_478)
        t_479_padded = F.pad(t_479, [1, 1, 1, 1], value=0)
        t_480 = self.n_Conv_53(t_479_padded)
        t_481 = F.relu(t_480)
        t_482 = self.n_Conv_54(t_481)
        t_483 = torch.add(t_482, t_477)
        t_484 = F.relu(t_483)
        t_485 = self.n_Conv_55(t_484)
        t_486 = F.relu(t_485)
        t_486_padded = F.pad(t_486, [1, 1, 1, 1], value=0)
        t_487 = self.n_Conv_56(t_486_padded)
        t_488 = F.relu(t_487)
        t_489 = self.n_Conv_57(t_488)
        t_490 = torch.add(t_489, t_484)
        t_491 = F.relu(t_490)
        t_492 = self.n_Conv_58(t_491)
        t_493 = F.relu(t_492)
        t_493_padded = F.pad(t_493, [1, 1, 1, 1], value=0)
        t_494 = self.n_Conv_59(t_493_padded)
        t_495 = F.relu(t_494)
        t_496 = self.n_Conv_60(t_495)
        t_497 = torch.add(t_496, t_491)
        t_498 = F.relu(t_497)
        t_499 = self.n_Conv_61(t_498)
        t_500 = F.relu(t_499)
        t_500_padded = F.pad(t_500, [1, 1, 1, 1], value=0)
        t_501 = self.n_Conv_62(t_500_padded)
        t_502 = F.relu(t_501)
        t_503 = self.n_Conv_63(t_502)
        t_504 = torch.add(t_503, t_498)
        t_505 = F.relu(t_504)
        t_506 = self.n_Conv_64(t_505)
        t_507 = F.relu(t_506)
        t_507_padded = F.pad(t_507, [1, 1, 1, 1], value=0)
        t_508 = self.n_Conv_65(t_507_padded)
        t_509 = F.relu(t_508)
        t_510 = self.n_Conv_66(t_509)
        t_511 = torch.add(t_510, t_505)
        t_512 = F.relu(t_511)
        t_513 = self.n_Conv_67(t_512)
        t_514 = F.relu(t_513)
        t_514_padded = F.pad(t_514, [1, 1, 1, 1], value=0)
        t_515 = self.n_Conv_68(t_514_padded)
        t_516 = F.relu(t_515)
        t_517 = self.n_Conv_69(t_516)
        t_518 = torch.add(t_517, t_512)
        t_519 = F.relu(t_518)
        t_520 = self.n_Conv_70(t_519)
        t_521 = F.relu(t_520)
        t_521_padded = F.pad(t_521, [1, 1, 1, 1], value=0)
        t_522 = self.n_Conv_71(t_521_padded)
        t_523 = F.relu(t_522)
        t_524 = self.n_Conv_72(t_523)
        t_525 = torch.add(t_524, t_519)
        t_526 = F.relu(t_525)
        t_527 = self.n_Conv_73(t_526)
        t_528 = F.relu(t_527)
        t_528_padded = F.pad(t_528, [1, 1, 1, 1], value=0)
        t_529 = self.n_Conv_74(t_528_padded)
        t_530 = F.relu(t_529)
        t_531 = self.n_Conv_75(t_530)
        t_532 = torch.add(t_531, t_526)
        t_533 = F.relu(t_532)
        t_534 = self.n_Conv_76(t_533)
        t_535 = F.relu(t_534)
        t_535_padded = F.pad(t_535, [1, 1, 1, 1], value=0)
        t_536 = self.n_Conv_77(t_535_padded)
        t_537 = F.relu(t_536)
        t_538 = self.n_Conv_78(t_537)
        t_539 = torch.add(t_538, t_533)
        t_540 = F.relu(t_539)
        t_541 = self.n_Conv_79(t_540)
        t_542 = F.relu(t_541)
        t_542_padded = F.pad(t_542, [1, 1, 1, 1], value=0)
        t_543 = self.n_Conv_80(t_542_padded)
        t_544 = F.relu(t_543)
        t_545 = self.n_Conv_81(t_544)
        t_546 = torch.add(t_545, t_540)
        t_547 = F.relu(t_546)
        t_548 = self.n_Conv_82(t_547)
        t_549 = F.relu(t_548)
        t_549_padded = F.pad(t_549, [1, 1, 1, 1], value=0)
        t_550 = self.n_Conv_83(t_549_padded)
        t_551 = F.relu(t_550)
        t_552 = self.n_Conv_84(t_551)
        t_553 = torch.add(t_552, t_547)
        t_554 = F.relu(t_553)
        t_555 = self.n_Conv_85(t_554)
        t_556 = F.relu(t_555)
        t_556_padded = F.pad(t_556, [1, 1, 1, 1], value=0)
        t_557 = self.n_Conv_86(t_556_padded)
        t_558 = F.relu(t_557)
        t_559 = self.n_Conv_87(t_558)
        t_560 = torch.add(t_559, t_554)
        t_561 = F.relu(t_560)
        t_562 = self.n_Conv_88(t_561)
        t_563 = F.relu(t_562)
        t_563_padded = F.pad(t_563, [1, 1, 1, 1], value=0)
        t_564 = self.n_Conv_89(t_563_padded)
        t_565 = F.relu(t_564)
        t_566 = self.n_Conv_90(t_565)
        t_567 = torch.add(t_566, t_561)
        t_568 = F.relu(t_567)
        t_569 = self.n_Conv_91(t_568)
        t_570 = F.relu(t_569)
        t_570_padded = F.pad(t_570, [1, 1, 1, 1], value=0)
        t_571 = self.n_Conv_92(t_570_padded)
        t_572 = F.relu(t_571)
        t_573 = self.n_Conv_93(t_572)
        t_574 = torch.add(t_573, t_568)
        t_575 = F.relu(t_574)
        t_576 = self.n_Conv_94(t_575)
        t_577 = F.relu(t_576)
        t_577_padded = F.pad(t_577, [1, 1, 1, 1], value=0)
        t_578 = self.n_Conv_95(t_577_padded)
        t_579 = F.relu(t_578)
        t_580 = self.n_Conv_96(t_579)
        t_581 = torch.add(t_580, t_575)
        t_582 = F.relu(t_581)
        t_583 = self.n_Conv_97(t_582)
        t_584 = F.relu(t_583)
        t_584_padded = F.pad(t_584, [0, 1, 0, 1], value=0)
        t_585 = self.n_Conv_98(t_584_padded)
        t_586 = F.relu(t_585)
        t_587 = self.n_Conv_99(t_586)
        t_588 = self.n_Conv_100(t_582)
        t_589 = torch.add(t_587, t_588)
        t_590 = F.relu(t_589)
        t_591 = self.n_Conv_101(t_590)
        t_592 = F.relu(t_591)
        t_592_padded = F.pad(t_592, [1, 1, 1, 1], value=0)
        t_593 = self.n_Conv_102(t_592_padded)
        t_594 = F.relu(t_593)
        t_595 = self.n_Conv_103(t_594)
        t_596 = torch.add(t_595, t_590)
        t_597 = F.relu(t_596)
        t_598 = self.n_Conv_104(t_597)
        t_599 = F.relu(t_598)
        t_599_padded = F.pad(t_599, [1, 1, 1, 1], value=0)
        t_600 = self.n_Conv_105(t_599_padded)
        t_601 = F.relu(t_600)
        t_602 = self.n_Conv_106(t_601)
        t_603 = torch.add(t_602, t_597)
        t_604 = F.relu(t_603)
        t_605 = self.n_Conv_107(t_604)
        t_606 = F.relu(t_605)
        t_606_padded = F.pad(t_606, [1, 1, 1, 1], value=0)
        t_607 = self.n_Conv_108(t_606_padded)
        t_608 = F.relu(t_607)
        t_609 = self.n_Conv_109(t_608)
        t_610 = torch.add(t_609, t_604)
        t_611 = F.relu(t_610)
        t_612 = self.n_Conv_110(t_611)
        t_613 = F.relu(t_612)
        t_613_padded = F.pad(t_613, [1, 1, 1, 1], value=0)
        t_614 = self.n_Conv_111(t_613_padded)
        t_615 = F.relu(t_614)
        t_616 = self.n_Conv_112(t_615)
        t_617 = torch.add(t_616, t_611)
        t_618 = F.relu(t_617)
        t_619 = self.n_Conv_113(t_618)
        t_620 = F.relu(t_619)
        t_620_padded = F.pad(t_620, [1, 1, 1, 1], value=0)
        t_621 = self.n_Conv_114(t_620_padded)
        t_622 = F.relu(t_621)
        t_623 = self.n_Conv_115(t_622)
        t_624 = torch.add(t_623, t_618)
        t_625 = F.relu(t_624)
        t_626 = self.n_Conv_116(t_625)
        t_627 = F.relu(t_626)
        t_627_padded = F.pad(t_627, [1, 1, 1, 1], value=0)
        t_628 = self.n_Conv_117(t_627_padded)
        t_629 = F.relu(t_628)
        t_630 = self.n_Conv_118(t_629)
        t_631 = torch.add(t_630, t_625)
        t_632 = F.relu(t_631)
        t_633 = self.n_Conv_119(t_632)
        t_634 = F.relu(t_633)
        t_634_padded = F.pad(t_634, [1, 1, 1, 1], value=0)
        t_635 = self.n_Conv_120(t_634_padded)
        t_636 = F.relu(t_635)
        t_637 = self.n_Conv_121(t_636)
        t_638 = torch.add(t_637, t_632)
        t_639 = F.relu(t_638)
        t_640 = self.n_Conv_122(t_639)
        t_641 = F.relu(t_640)
        t_641_padded = F.pad(t_641, [1, 1, 1, 1], value=0)
        t_642 = self.n_Conv_123(t_641_padded)
        t_643 = F.relu(t_642)
        t_644 = self.n_Conv_124(t_643)
        t_645 = torch.add(t_644, t_639)
        t_646 = F.relu(t_645)
        t_647 = self.n_Conv_125(t_646)
        t_648 = F.relu(t_647)
        t_648_padded = F.pad(t_648, [1, 1, 1, 1], value=0)
        t_649 = self.n_Conv_126(t_648_padded)
        t_650 = F.relu(t_649)
        t_651 = self.n_Conv_127(t_650)
        t_652 = torch.add(t_651, t_646)
        t_653 = F.relu(t_652)
        t_654 = self.n_Conv_128(t_653)
        t_655 = F.relu(t_654)
        t_655_padded = F.pad(t_655, [1, 1, 1, 1], value=0)
        t_656 = self.n_Conv_129(t_655_padded)
        t_657 = F.relu(t_656)
        t_658 = self.n_Conv_130(t_657)
        t_659 = torch.add(t_658, t_653)
        t_660 = F.relu(t_659)
        t_661 = self.n_Conv_131(t_660)
        t_662 = F.relu(t_661)
        t_662_padded = F.pad(t_662, [1, 1, 1, 1], value=0)
        t_663 = self.n_Conv_132(t_662_padded)
        t_664 = F.relu(t_663)
        t_665 = self.n_Conv_133(t_664)
        t_666 = torch.add(t_665, t_660)
        t_667 = F.relu(t_666)
        t_668 = self.n_Conv_134(t_667)
        t_669 = F.relu(t_668)
        t_669_padded = F.pad(t_669, [1, 1, 1, 1], value=0)
        t_670 = self.n_Conv_135(t_669_padded)
        t_671 = F.relu(t_670)
        t_672 = self.n_Conv_136(t_671)
        t_673 = torch.add(t_672, t_667)
        t_674 = F.relu(t_673)
        t_675 = self.n_Conv_137(t_674)
        t_676 = F.relu(t_675)
        t_676_padded = F.pad(t_676, [1, 1, 1, 1], value=0)
        t_677 = self.n_Conv_138(t_676_padded)
        t_678 = F.relu(t_677)
        t_679 = self.n_Conv_139(t_678)
        t_680 = torch.add(t_679, t_674)
        t_681 = F.relu(t_680)
        t_682 = self.n_Conv_140(t_681)
        t_683 = F.relu(t_682)
        t_683_padded = F.pad(t_683, [1, 1, 1, 1], value=0)
        t_684 = self.n_Conv_141(t_683_padded)
        t_685 = F.relu(t_684)
        t_686 = self.n_Conv_142(t_685)
        t_687 = torch.add(t_686, t_681)
        t_688 = F.relu(t_687)
        t_689 = self.n_Conv_143(t_688)
        t_690 = F.relu(t_689)
        t_690_padded = F.pad(t_690, [1, 1, 1, 1], value=0)
        t_691 = self.n_Conv_144(t_690_padded)
        t_692 = F.relu(t_691)
        t_693 = self.n_Conv_145(t_692)
        t_694 = torch.add(t_693, t_688)
        t_695 = F.relu(t_694)
        t_696 = self.n_Conv_146(t_695)
        t_697 = F.relu(t_696)
        t_697_padded = F.pad(t_697, [1, 1, 1, 1], value=0)
        t_698 = self.n_Conv_147(t_697_padded)
        t_699 = F.relu(t_698)
        t_700 = self.n_Conv_148(t_699)
        t_701 = torch.add(t_700, t_695)
        t_702 = F.relu(t_701)
        t_703 = self.n_Conv_149(t_702)
        t_704 = F.relu(t_703)
        t_704_padded = F.pad(t_704, [1, 1, 1, 1], value=0)
        t_705 = self.n_Conv_150(t_704_padded)
        t_706 = F.relu(t_705)
        t_707 = self.n_Conv_151(t_706)
        t_708 = torch.add(t_707, t_702)
        t_709 = F.relu(t_708)
        t_710 = self.n_Conv_152(t_709)
        t_711 = F.relu(t_710)
        t_711_padded = F.pad(t_711, [1, 1, 1, 1], value=0)
        t_712 = self.n_Conv_153(t_711_padded)
        t_713 = F.relu(t_712)
        t_714 = self.n_Conv_154(t_713)
        t_715 = torch.add(t_714, t_709)
        t_716 = F.relu(t_715)
        t_717 = self.n_Conv_155(t_716)
        t_718 = F.relu(t_717)
        t_718_padded = F.pad(t_718, [1, 1, 1, 1], value=0)
        t_719 = self.n_Conv_156(t_718_padded)
        t_720 = F.relu(t_719)
        t_721 = self.n_Conv_157(t_720)
        t_722 = torch.add(t_721, t_716)
        t_723 = F.relu(t_722)
        t_724 = self.n_Conv_158(t_723)
        t_725 = self.n_Conv_159(t_723)
        t_726 = F.relu(t_725)
        t_726_padded = F.pad(t_726, [0, 1, 0, 1], value=0)
        t_727 = self.n_Conv_160(t_726_padded)
        t_728 = F.relu(t_727)
        t_729 = self.n_Conv_161(t_728)
        t_730 = torch.add(t_729, t_724)
        t_731 = F.relu(t_730)
        t_732 = self.n_Conv_162(t_731)
        t_733 = F.relu(t_732)
        t_733_padded = F.pad(t_733, [1, 1, 1, 1], value=0)
        t_734 = self.n_Conv_163(t_733_padded)
        t_735 = F.relu(t_734)
        t_736 = self.n_Conv_164(t_735)
        t_737 = torch.add(t_736, t_731)
        t_738 = F.relu(t_737)
        t_739 = self.n_Conv_165(t_738)
        t_740 = F.relu(t_739)
        t_740_padded = F.pad(t_740, [1, 1, 1, 1], value=0)
        t_741 = self.n_Conv_166(t_740_padded)
        t_742 = F.relu(t_741)
        t_743 = self.n_Conv_167(t_742)
        t_744 = torch.add(t_743, t_738)
        t_745 = F.relu(t_744)
        t_746 = self.n_Conv_168(t_745)
        t_747 = self.n_Conv_169(t_745)
        t_748 = F.relu(t_747)
        t_748_padded = F.pad(t_748, [0, 1, 0, 1], value=0)
        t_749 = self.n_Conv_170(t_748_padded)
        t_750 = F.relu(t_749)
        t_751 = self.n_Conv_171(t_750)
        t_752 = torch.add(t_751, t_746)
        t_753 = F.relu(t_752)
        t_754 = self.n_Conv_172(t_753)
        t_755 = F.relu(t_754)
        t_755_padded = F.pad(t_755, [1, 1, 1, 1], value=0)
        t_756 = self.n_Conv_173(t_755_padded)
        t_757 = F.relu(t_756)
        t_758 = self.n_Conv_174(t_757)
        t_759 = torch.add(t_758, t_753)
        t_760 = F.relu(t_759)
        t_761 = self.n_Conv_175(t_760)
        t_762 = F.relu(t_761)
        t_762_padded = F.pad(t_762, [1, 1, 1, 1], value=0)
        t_763 = self.n_Conv_176(t_762_padded)
        t_764 = F.relu(t_763)
        t_765 = self.n_Conv_177(t_764)
        t_766 = torch.add(t_765, t_760)
        t_767 = F.relu(t_766)
        t_768 = self.n_Conv_178(t_767)
        t_769 = F.avg_pool2d(t_768, kernel_size=t_768.shape[-2:])
        t_770 = torch.squeeze(t_769, 3)
        t_770 = torch.squeeze(t_770, 2)
        t_771 = torch.sigmoid(t_770)
        return t_771

    def load_state_dict(self, state_dict, **kwargs):
        self.tags = state_dict.get("tags", [])

        super(DeepDanbooruModel, self).load_state_dict(
            {k: v for k, v in state_dict.items() if k != "tags"}
        )


class DeepDanbooruWrapper:
    def __init__(
        self,
        device,
        threshold=0.5,
        alpha_sort=False,
        use_spaces=True,
        use_escape=True,
        include_ranks=False,
    ):
        self.device = device
        self.model = None
        self.threshold = threshold
        self.alpha_sort = alpha_sort
        self.use_spaces = use_spaces
        self.use_escape = use_escape
        self.include_ranks = include_ranks

    def load_models(self, model_path):

        self.model = DeepDanbooruModel()
        state_dict = torch.load(model_path)
        self.model.load_state_dict(state_dict, map_location="cpu")

        self.model.eval()

    def start(self):
        
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()
        gc.collect()
        assert self.model is not None, "please load deepbooru model first"
        self.model.to(self.device)

    def stop(self):
        self.model.to("cpu")
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()
        gc.collect()

    def __call__(self, pil_image):
        self.start()
        res = self.tag_multi(pil_image)
        self.stop()
        return res

    def tag_multi(self, pil_image, force_disable_ranks=False):

        pil_image = pil_image.resize((512, 512)).convert("RGB")
        image_np = np.expand_dims(np.array(pil_image, dtype=np.float32), 0) / 255
        image_tensor = torch.from_numpy(image_np).to(torch.float32).to(self.device)
        with torch.no_grad():
            y = self.model(image_tensor)[0].detach().cpu().numpy()
        probability_dict = {}

        for tag, probability in zip(self.model.tags, y):
            if probability < self.threshold:
                continue

            if tag.startswith("rating:"):
                continue

            probability_dict[tag] = probability

        if self.alpha_sort:
            tags = sorted(probability_dict)
        else:
            tags = [
                tag for tag, _ in sorted(probability_dict.items(), key=lambda x: -x[1])
            ]

        res = []

        # filtertags = {x.strip().replace(' ', '_') for x in shared.opts.deepbooru_filter_tags.split(",")}
        filtertags = []
        for tag in [x for x in tags if x not in filtertags]:
            probability = probability_dict[tag]
            tag_outformat = tag
            if self.use_spaces:
                tag_outformat = tag_outformat.replace("_", " ")
            if self.use_escape:
                tag_outformat = re.sub(re_special, r"\\\1", tag_outformat)
            if self.include_ranks:
                tag_outformat = f"({tag_outformat}:{probability:.3f})"

            res.append(tag_outformat)

        return ", ".join(res)

modules/wrappers/deepbooru.py
MooreThreads/MobiMaliangSDK
import torch
import os
import sys
import gc
import re
from PIL import Image
from io import BytesIO
import base64
import numpy as np
from realesrgan import RealESRGANer
from basicsr.archs.rrdbnet_arch import RRDBNet

import os
import sys
__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.shared import MTGPU_DETECTION, device


class RealESRGANWrapper:
    def __init__(self, device=device):
        self.device = device
        self.model = None

    def load_models(self, model_path):
        RealESRGAN_models = {
            "RealESRGAN_x4plus": RRDBNet(
                num_in_ch=3,
                num_out_ch=3,
                num_feat=64,
                num_block=23,
                num_grow_ch=32,
                scale=4,
            )
        }

        RealESRGAN_dir = os.path.dirname(model_path)
        sys.path.append(os.path.abspath(RealESRGAN_dir))
        model_name = os.path.basename(model_path).split(".")[0]
        self.model = RealESRGANer(
            scale=2,
            model_path=model_path,
            model=RealESRGAN_models[model_name],
            pre_pad=0,
            half=False,
            device="cpu",
        )
        self.model.model.name = model_name

    def start(self):
        
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()
        gc.collect()
        assert self.model is not None, "please load deepbooru model first"
        self.model.model.to(self.device)
        self.model.device = self.device

    def stop(self):
        
        self.model.model.to("cpu")
        self.model.device = "cpu"
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()
        gc.collect()

    def super_resolution(self, image):
        assert image is not None, "Input image should not be None"
        image = image.convert("RGB")
        np_image = np.array(image, dtype=np.uint8)
        output, _ = self.model.enhance(np_image)
        res = Image.fromarray(output)
        return res

    def __call__(self, image):

        self.start()
        res = self.super_resolution(image)
        self.stop()

        return res

modules/wrappers/realesrgan.py
MooreThreads/MobiMaliangSDK
import torch
import gc
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

import os
import sys
__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(__dir__))
from modules.shared import MTGPU_DETECTION, device

class Translator:
    def __init__(self, device=device):
        self.device = device
        self.tokenizer = None
        self.model = None

    def load_models(self, model_path):
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path)

    def start(self):
        
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()
        gc.collect()
        assert self.model is not None, "please load translator model first"
        self.model.to(self.device)

    def stop(self):
        
        self.model.to("cpu")
        torch.musa.empty_cache() if MTGPU_DETECTION else torch.cuda.empty_cache()
        gc.collect()

    def __call__(self, text):

        self.start()

        tokenized_text = self.tokenizer([text], return_tensors="pt")
        translation = self.model.generate(**tokenized_text)

        self.stop()

        return self.tokenizer.batch_decode(translation, skip_special_tokens=True)[0]

modules/wrappers/translator.py
MooreThreads/MobiMaliangSDK