new start

.gitignore

@@ -0,0 +1 @@
+venv/*

anime_app.py

@@ -1,135 +1,136 @@
-import gradio as gr
-from anime_model import Model
-import spaces
-
-show_options = True
-    
-from settings import (
-    DEFAULT_IMAGE_RESOLUTION,
-    MAX_NUM_IMAGES,
-    MAX_SEED,
-)
-from utils import randomize_seed_fn
-# if gr.NO_RELOAD:
-
-@spaces.GPU
-def init():
-    base_model = "nyxia/AAM-AnyLoRA-Anime-Mix"
-    model = Model()
-
-def auto_process_image(image, prompt):
-    a_prompt="anime style, cartoon, drawing, 2D anime, illustration, cartoon"
-    n_prompt="realism, 3d,  BadDream, (UnrealisticDream:1.2), split image, multiple views, text, cropped, out of frame, worst quality, low quality, jpeg artifacts, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, bad proportions, extra limbs, longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
-    num_samples=2
-    image_resolution=512
-    preprocess_resolution=512
-    num_steps=15
-    guidance_scale=4.5
-    seed=randomize_seed_fn(0, True)
-    config = [
-        image,
-        prompt,
-        a_prompt,
-        n_prompt,
-        num_samples,
-        image_resolution,
-        preprocess_resolution,
-        num_steps,
-        guidance_scale,
-        seed,
-    ]
-    if image is None:
-        return None
-    print("processing image")
-    config[0] = image
-    config[1] = prompt
-    # print(config)
-    return model.process_normal(*config)    
-
-with gr.Blocks(theme=gr.themes.Soft(), css="footer {visibility: hidden}") as demo:
-    # with gr.Row():
-    #     # examples
-    #     gr.Text(label="Anime Style Examples", value="Weeb!")
-    with gr.Row():
-        with gr.Column():
-            # input text
-            prompt = gr.Textbox(label="Anime Style", placeholder="anime tittes")
-    with gr.Row():
-        with gr.Column():
-            # input image
-            image = gr.Image(label="Input", sources=['upload'], show_label=True, format="jpeg")
-        with gr.Column():
-            # output
-            result = gr.Gallery(label="Anime", show_label=True, columns=2, scale=3, object_fit="contain", format="jpeg")
-    with gr.Column():
-            # run button
-            run_button = gr.Button(size=["lg"])
-    with gr.Row():
-        with gr.Accordion("Advanced options", open=show_options, visible=show_options):
-            num_samples = gr.Slider(
-                label="Images", minimum=1, maximum=MAX_NUM_IMAGES, value=2, step=1
-            )
-            image_resolution = gr.Slider(
-                label="Image resolution",
-                minimum=256,
-                maximum=1024,
-                value=DEFAULT_IMAGE_RESOLUTION,
-                step=256,
-            )
-            preprocess_resolution = gr.Slider(
-                label="Preprocess resolution", minimum=128, maximum=1024, value=512, step=1
-            )
-            num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=15, step=1) # 20/4.5 or 12 without lora, 4 with lora
-            guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=4.5, step=0.1) #5 without lora, 2 with lora
-            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            a_prompt = gr.Textbox(label="Additional prompt", value="anime style, cartoon, drawing, 2D anime, illustration, cartoon")
-            n_prompt = gr.Textbox(
-                label="Negative prompt",
-                # value="BadDream, (UnrealisticDream:1.2), split image, multiple views, text, cropped, out of frame, worst quality, low quality, jpeg artifacts, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, bad proportions, extra limbs, longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
-                value="(signature:1.2),(artist name:1.2),(watermark:1.2), (easynegative), (low quality, worst quality:1.4), (bad anatomy), (inaccurate limb:1.2),bad composition, inaccurate eyes, extra digit,fewer digits, (extra arms:1.2), badhandv4,clothes"
-            )
-    config = [
-        image,
-        prompt,
-        a_prompt,
-        n_prompt,
-        num_samples,
-        image_resolution,
-        preprocess_resolution,
-        num_steps,
-        guidance_scale,
-        seed,
-    ]
-    prompt.submit(
-        fn=randomize_seed_fn,
-        inputs=[seed, randomize_seed],
-        outputs=seed,
-        queue=False,
-        api_name=False,
-        show_progress="minimal",
-    ).then(
-        fn=Model.process_normal,
-        inputs=config,
-        outputs=result,
-        api_name=False,
-        show_progress="minimal"
-    )
-    run_button.click(
-        fn=randomize_seed_fn,
-        inputs=[seed, randomize_seed],
-        outputs=seed,
-        queue=False,
-        api_name=False,
-        show_progress="minimal"
-    ).then(
-        fn=Model.process_normal,
-        inputs=config,
-        outputs=result,
-        show_progress="minimal"
-    )
-#    image.change(auto_process_image, inputs=[image, prompt], outputs=[result])
-
-if __name__ == "__main__":
-    #init()
-    demo.queue(max_size=1).launch()
+import gradio as gr
+from anime_model import Model
+import spaces
+prod = False
+port = 8080
+show_options = True
+if prod:
+    port = 8081
+    show_options = False
+    
+from settings import (
+    DEFAULT_IMAGE_RESOLUTION,
+    MAX_NUM_IMAGES,
+    MAX_SEED,
+)
+from utils import randomize_seed_fn
+
+base_model = "nyxia/AAM-AnyLoRA-Anime-Mix"
+model = Model(base_model_id=base_model, task_name="NormalBae")
+
+# note: for high res 1024x1024, set num steps to 9 and guidance to 6
+def auto_process_image(image, prompt):
+    a_prompt="anime style, cartoon, drawing, 2D anime, illustration, cartoon"
+    n_prompt="realism, 3d,  BadDream, (UnrealisticDream:1.2), split image, multiple views, text, cropped, out of frame, worst quality, low quality, jpeg artifacts, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, bad proportions, extra limbs, longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
+    num_samples=4
+    image_resolution=512
+    preprocess_resolution=512
+    num_steps=15
+    guidance_scale=4.5
+    seed=randomize_seed_fn(0, True)
+    config = [
+        image,
+        prompt,
+        a_prompt,
+        n_prompt,
+        num_samples,
+        image_resolution,
+        preprocess_resolution,
+        num_steps,
+        guidance_scale,
+        seed,
+    ]
+    if image is None:
+        return None
+    print("processing image")
+    config[0] = image
+    config[1] = prompt
+    # print(config)
+    return model.process_normal(*config)    
+
+with gr.Blocks(theme=gr.themes.Soft(), css="footer {visibility: hidden}") as demo:
+    with gr.Row():
+        # examples
+        gr.Text(label="Anime Style Examples", value="Weeb!")
+    with gr.Row():
+        with gr.Column():
+            # input text
+            prompt = gr.Textbox(label="Anime Style", placeholder="anime tittes")
+    with gr.Row():
+        with gr.Column():
+            # input image
+            image = gr.Image(label="Input", sources=['upload'], show_label=True, format="jpeg")
+        with gr.Column():
+            # output
+            result = gr.Gallery(label="Anime", show_label=True, columns=2, scale=3, object_fit="contain", format="jpeg")
+    with gr.Column():
+            # run button
+            run_button = gr.Button(size=["lg"])
+    with gr.Row():
+        with gr.Accordion("Advanced options", open=show_options, visible=show_options):
+            num_samples = gr.Slider(
+                label="Images", minimum=1, maximum=MAX_NUM_IMAGES, value=4, step=1
+            )
+            image_resolution = gr.Slider(
+                label="Image resolution",
+                minimum=256,
+                maximum=1024,
+                value=DEFAULT_IMAGE_RESOLUTION,
+                step=256,
+            )
+            preprocess_resolution = gr.Slider(
+                label="Preprocess resolution", minimum=128, maximum=1024, value=512, step=1
+            )
+            num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=15, step=1) # 20/4.5 or 12 without lora, 4 with lora
+            guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=4.5, step=0.1) #5 without lora, 2 with lora
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            a_prompt = gr.Textbox(label="Additional prompt", value="anime style, cartoon, drawing, 2D anime, illustration, cartoon")
+            n_prompt = gr.Textbox(
+                label="Negative prompt",
+                # value="BadDream, (UnrealisticDream:1.2), split image, multiple views, text, cropped, out of frame, worst quality, low quality, jpeg artifacts, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, bad proportions, extra limbs, longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
+                value="(signature:1.2),(artist name:1.2),(watermark:1.2), (easynegative), (low quality, worst quality:1.4), (bad anatomy), (inaccurate limb:1.2),bad composition, inaccurate eyes, extra digit,fewer digits, (extra arms:1.2), badhandv4,clothes"
+            )
+    config = [
+        image,
+        prompt,
+        a_prompt,
+        n_prompt,
+        num_samples,
+        image_resolution,
+        preprocess_resolution,
+        num_steps,
+        guidance_scale,
+        seed,
+    ]
+    prompt.submit(
+        fn=randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+        show_progress="minimal",
+    ).then(
+        fn=model.process_normal,
+        inputs=config,
+        outputs=result,
+        api_name=False,
+        show_progress="minimal"
+    )
+    run_button.click(
+        fn=randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+        show_progress="minimal"
+    ).then(
+        fn=model.process_normal,
+        inputs=config,
+        outputs=result,
+        show_progress="minimal"
+    )
+    image.change(auto_process_image, inputs=[image, prompt], outputs=[result])
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch(server_name="localhost", server_port=port)

anime_model.py

@@ -1,185 +1,157 @@
-from __future__ import annotations
-import gc
-import time
-import spaces
-import numpy as np
-import PIL.Image
-import torch
-from controlnet_aux import NormalBaeDetector
-from controlnet_aux.util import HWC3
-from cv_utils import resize_image
-from diffusers import (
-    ControlNetModel,
-    AutoencoderKL,
-    StableDiffusionControlNetPipeline,
-    DPMSolverMultistepScheduler,
-)
-from settings import MAX_IMAGE_RESOLUTION, MAX_NUM_IMAGES
-
-class Model:
-
-    def __init__(self):
-        print("Initializing base model: ", base_model_id, " with ", task_name)
-        self.base_model_id = "nyxia/AAM-AnyLoRA-Anime-Mix"
-        self.task_name = "NormalBae"
-        self.pipe = self.load_pipe(base_model_id, task_name)
-        self.model = None
-        self.name = ""
-
-    def load(self, name: str) -> None:
-        if name == self.name:
-            return
-        elif name == "NormalBae":
-            model = NormalBaeDetector.from_pretrained(self.MODEL_ID)#.to("cuda")
-        else:
-            raise ValueError
-        torch.cuda.empty_cache()
-        gc.collect()
-        self.name = name
-        self.model = model
-
-
-    def preprocess(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
-        if self.name == "Canny":
-            if "detect_resolution" in kwargs:
-                detect_resolution = kwargs.pop("detect_resolution")
-                image = np.array(image)
-                image = HWC3(image)
-                image = resize_image(image, resolution=detect_resolution)
-            image = self.model(image, **kwargs)
-            return PIL.Image.fromarray(image)
-        elif self.name == "Midas":
-            detect_resolution = kwargs.pop("detect_resolution", 512)
-            image_resolution = kwargs.pop("image_resolution", 512)
-            image = np.array(image)
-            image = HWC3(image)
-            image = resize_image(image, resolution=detect_resolution)
-            image = self.model(image, **kwargs)
-            image = HWC3(image)
-            image = resize_image(image, resolution=image_resolution)
-            return PIL.Image.fromarray(image)
-        else:
-            return self.model(image, **kwargs)
-    
-    @spaces.GPU
-    def load_pipe(self, base_model_id, task_name):
-        print("loading pipe")
-        # Controlnet
-        model_id = "lllyasviel/control_v11p_sd15_normalbae"
-        print("initializing controlnet")
-        controlnet = ControlNetModel.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16,
-            attn_implementation="flash_attention_2",
-        ).to("cuda")
-        controlnet.to(memory_format=torch.channels_last)
-        
-        # Scheduler
-        scheduler = DPMSolverMultistepScheduler.from_pretrained(
-            "stabilityai/stable-diffusion-xl-base-1.0",
-            subfolder="scheduler",
-            use_karras_sigmas=True,
-            algorithm_type="sde-dpmsolver++",
-            denoise_final=True,
-            device_map="cuda",
-            attn_implementation="flash_attention_2",
-        )
-
-        # VAE
-        vae_url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors"
-        vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda")
-        vae.to(memory_format=torch.channels_last)
-
-        # Stable Diffusion Pipeline
-        pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            base_model_id,
-            safety_checker=None,
-            controlnet=controlnet,
-            scheduler=scheduler,
-            vae=vae,
-            torch_dtype=torch.float16,
-        ).to("cuda")
-
-        # Efficiency optimizations - DO NOT CHANGE ORDER
-        pipe.enable_xformers_memory_efficient_attention()
-
-        torch.cuda.empty_cache()
-        gc.collect()
-        self.base_model_id = base_model_id
-        self.task_name = task_name
-
-        return pipe
-
-    def get_prompt(self, prompt: str, additional_prompt: str) -> str:
-        if not prompt:
-            prompt = additional_prompt
-        else:
-            prompt = f"{prompt}, {additional_prompt}"
-        return prompt
-
-    @torch.inference_mode()
-    def run_pipe(
-        prompt: str,
-        negative_prompt: str,
-        control_image: PIL.Image.Image,
-        num_images: int,
-        num_steps: int,
-        guidance_scale: float,
-        seed: int,
-    ) -> list[PIL.Image.Image]:
-        generator = torch.cuda.manual_seed(seed)
-        torch.cuda.synchronize()
-        start = time.time()
-        results = self.pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            guidance_scale=guidance_scale,
-            num_images_per_prompt=num_images,
-            num_inference_steps=num_steps,
-            generator=generator,
-            image=control_image,
-        ).images
-        print(f"Inference done in: {time.time() - start:.2f} seconds")
-        print(f"Prompt {prompt}")
-        torch.cuda.synchronize()
-        torch.cuda.empty_cache()
-        gc.collect()
-
-        return results
-
-    @spaces.GPU
-    def process_normal(
-        self,
-        image: np.ndarray,
-        prompt: str,
-        additional_prompt: str,
-        negative_prompt: str,
-        num_images: int,
-        image_resolution: int,
-        preprocess_resolution: int,
-        num_steps: int,
-        guidance_scale: float,
-        seed: int,
-    ) -> list[PIL.Image.Image]:
-        if image is None:
-            raise ValueError
-        #if image_resolution > MAX_IMAGE_RESOLUTION:
-        #    raise ValueError
-        #if num_images > MAX_NUM_IMAGES:
-        #    raise ValueError
-        #self.load("NormalBae")
-        model = NormalBaeDetector.from_pretrained("lllyasviel/Annotators").to("cuda")
-        torch.cuda.empty_cache()
-        gc.collect()
-        if prompt == "":
-            prompt = "anime girl"
-        print(prompt)
-        return run_pipe(
-            prompt=self.get_prompt("Hentai Nude Anime Titties of " + prompt, additional_prompt),
-            negative_prompt=negative_prompt,
-            control_image=image,
-            num_images=num_images,
-            num_steps=num_steps,
-            guidance_scale=guidance_scale,
-            seed=seed,
-        )
+from __future__ import annotations
+
+import gc
+import time
+import numpy as np
+import PIL.Image
+import torch
+import spaces
+from diffusers import (
+    ControlNetModel,
+    AutoencoderKL,
+    StableDiffusionControlNetPipeline,
+    DPMSolverMultistepScheduler,
+) 
+from preprocess_anime import Preprocessor
+from settings import MAX_IMAGE_RESOLUTION, MAX_NUM_IMAGES
+
+@spaces.GPU
+class Model:
+    def __init__(self, base_model_id, task_name):
+        print("Initializing base model: ", base_model_id, " with ", task_name)
+        self.base_model_id = base_model_id
+        self.task_name = task_name
+        self.pipe = self.load_pipe(base_model_id, task_name)
+        self.preprocessor = Preprocessor()
+        
+    def load_pipe(self, base_model_id, task_name):
+        print("loading pipe")
+        # Controlnet
+        model_id = "lllyasviel/control_v11p_sd15_normalbae"
+        print("initializing controlnet")
+        controlnet = ControlNetModel.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16,
+            attn_implementation="flash_attention_2",            
+            ).to("cuda")
+        controlnet.to(memory_format=torch.channels_last)
+        
+        # Scheduler
+        scheduler = DPMSolverMultistepScheduler.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0", 
+            subfolder="scheduler",
+            use_karras_sigmas=True,
+            # final_sigmas_type="sigma_min",
+            algorithm_type="sde-dpmsolver++",
+            # prediction_type="epsilon",
+            # thresholding=False,
+            denoise_final=True,
+            device_map="cuda",
+            attn_implementation="flash_attention_2",   
+        )
+
+        #vae
+        vae_url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors"
+        vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda")
+        vae.to(memory_format=torch.channels_last) 
+        # Stable Diffusion Pipeline
+        pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            base_model_id,
+            safety_checker=None,
+            controlnet=controlnet,
+            scheduler=scheduler,
+            vae=vae,
+            torch_dtype=torch.float16,
+        ).to("cuda")
+        
+        # efficiency optimizations - DO NOT CHANGE ORDER
+        pipe.enable_xformers_memory_efficient_attention()
+        
+        # lora
+        # pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
+        # pipe.load_lora_weights("Lykon/AnyLoRA", weight_name="AnyLoRA_bakedVae_blessed_fp16.safetensors")
+        # pipe.load_lora_weights("Lykon/AnyLoRA", weight_name="AnyLoRA_noVae_fp16-pruned.safetensors")
+        # pipe.fuse_lora() 
+        # pipe.unet.to(memory_format=torch.channels_last)
+                        
+        torch.cuda.empty_cache()
+        gc.collect()
+        self.base_model_id = base_model_id
+        self.task_name = task_name
+
+        return pipe
+
+    def get_prompt(self, prompt: str, additional_prompt: str) -> str:
+        if not prompt:
+            prompt = additional_prompt
+        else:
+            prompt = f"{prompt}, {additional_prompt}"
+        return prompt
+    
+    @torch.inference_mode()
+    def run_pipe(
+        self,
+        prompt: str,
+        negative_prompt: str,
+        control_image: PIL.Image.Image,
+        num_images: int,
+        num_steps: int,
+        guidance_scale: float,
+        seed: int,
+    ) -> list[PIL.Image.Image]:
+        generator = torch.cuda.manual_seed(seed)
+        torch.cuda.synchronize()
+        start = time.time()
+        results = self.pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=num_images,
+            num_inference_steps=num_steps,
+            generator=generator,
+            image=control_image,
+        ).images
+        print(f"Inference done in: {time.time() - start:.2f} seconds")
+        print(f"Prompt {prompt}") 
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        return results
+    def process_normal(
+        self,
+        image: np.ndarray,
+        prompt: str,
+        additional_prompt: str,
+        negative_prompt: str,
+        num_images: int,
+        image_resolution: int,
+        preprocess_resolution: int,
+        num_steps: int,
+        guidance_scale: float,
+        seed: int,
+    ) -> list[PIL.Image.Image]:
+        if image is None:
+            raise ValueError
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+        self.preprocessor.load("NormalBae")
+        control_image = self.preprocessor(
+            image=image,
+            image_resolution=image_resolution,
+            detect_resolution=preprocess_resolution,
+        )
+        if prompt == "":
+            prompt = "anime girl"
+        print(prompt)
+        return self.run_pipe(
+            prompt=self.get_prompt("Hentai Photo from imgur of " + prompt, additional_prompt),
+            negative_prompt=negative_prompt,
+            control_image=control_image,
+            num_images=num_images,
+            num_steps=num_steps,
+            guidance_scale=guidance_scale,
+            seed=seed,
+        )

preprocess_anime.py

@@ -1,57 +1,49 @@
-import gc
-import spaces
-import numpy as np
-import PIL.Image
-import torch
-from controlnet_aux import (
-
-    NormalBaeDetector,
-
-)
-from controlnet_aux.util import HWC3
-
-from cv_utils import resize_image
-
-
-
-class Preprocessor:
-    MODEL_ID = "lllyasviel/Annotators"
-
-    def __init__(self):
-        self.model = None
-        self.name = ""
-
-    @spaces.GPU
-    def load(self, name: str) -> None:
-        if name == self.name:
-            return
-        elif name == "NormalBae":
-            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
-        else:
-            raise ValueError
-        torch.cuda.empty_cache()
-        gc.collect()
-        self.name = name
-
-    @spaces.GPU
-    def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
-        if self.name == "Canny":
-            if "detect_resolution" in kwargs:
-                detect_resolution = kwargs.pop("detect_resolution")
-                image = np.array(image)
-                image = HWC3(image)
-                image = resize_image(image, resolution=detect_resolution)
-            image = self.model(image, **kwargs)
-            return PIL.Image.fromarray(image)
-        elif self.name == "Midas":
-            detect_resolution = kwargs.pop("detect_resolution", 512)
-            image_resolution = kwargs.pop("image_resolution", 512)
-            image = np.array(image)
-            image = HWC3(image)
-            image = resize_image(image, resolution=detect_resolution)
-            image = self.model(image, **kwargs)
-            image = HWC3(image)
-            image = resize_image(image, resolution=image_resolution)
-            return PIL.Image.fromarray(image)
-        else:
-            return self.model(image, **kwargs)
+import gc
+
+import numpy as np
+import PIL.Image
+import torch
+from controlnet_aux import NormalBaeDetector
+
+from controlnet_aux.util import HWC3
+from cv_utils import resize_image
+
+class Preprocessor:
+    MODEL_ID = "lllyasviel/Annotators"
+
+    def __init__(self):
+        self.model = None
+        self.name = ""
+
+    def load(self, name: str) -> None:
+        if name == self.name:
+            return
+        elif name == "NormalBae":
+            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
+        else:
+            raise ValueError
+        torch.cuda.empty_cache()
+        gc.collect()
+        self.name = name
+
+    def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
+        if self.name == "Canny":
+            if "detect_resolution" in kwargs:
+                detect_resolution = kwargs.pop("detect_resolution")
+                image = np.array(image)
+                image = HWC3(image)
+                image = resize_image(image, resolution=detect_resolution)
+            image = self.model(image, **kwargs)
+            return PIL.Image.fromarray(image)
+        elif self.name == "Midas":
+            detect_resolution = kwargs.pop("detect_resolution", 512)
+            image_resolution = kwargs.pop("image_resolution", 512)
+            image = np.array(image)
+            image = HWC3(image)
+            image = resize_image(image, resolution=detect_resolution)
+            image = self.model(image, **kwargs)
+            image = HWC3(image)
+            image = resize_image(image, resolution=image_resolution)
+            return PIL.Image.fromarray(image)
+        else:
+            return self.model(image, **kwargs)

requirements.txt

@@ -1,6 +1,6 @@
-torch==2.1.2
-torchvision
-torchaudio
+#torch==2.1.2
+#torchvision
+#torchaudio
 diffusers==0.27.2
 einops==0.6.1
 gradio==4.26.0
@@ -10,6 +10,6 @@ mediapipe==0.10.1
 opencv-python-headless==4.8.0.74
 safetensors==0.4.2
 transformers==4.39.3
-xformers==0.0.23.post1
+#xformers==0.0.23.post1
 accelerate==0.29.1
 #controlnet_aux==0.0.7