Upload 5 files

app.py

@@ -1,154 +1,6 @@
-import gradio as gr
-import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+from ui.gradio_ui import ui
+from sd.sd_controller import Controller
+
+controller=Controller()
+
+ui(controller)

sd/sd_controller.py

@@ -0,0 +1,74 @@
+from sd.utils.utils import *
+from utils.utils import sketch_process, prompt_preprocess
+#from controlnet_aux.pidi import PidiNetDetector
+import spaces
+
+class Controller():
+
+    def __init__(self, 
+                 models_names=["cagliostrolab/animagine-xl-3.1",
+                                     "stabilityai/stable-diffusion-xl-base-1.0"],
+                 lora_path='sd/lora/lora.safetensors'):
+        
+        self.models_names=models_names
+        self.lora_path=lora_path
+        self.vae=get_vae()
+        self.controlnet=get_controlnet()
+        self.adaptr=get_adapter()
+        self.scheduler=get_scheduler(model_name=self.models_names[1])
+        self.detector=get_detector()
+
+        self.first_pipe=get_pipe(vae=self.vae, 
+                                 model_name=self.models_names[0], 
+                                 controlnet=self.controlnet
+                                 lora_path=self.lora_path)
+
+        self.second_pipe=get_pipe(vae=self.vae, 
+                                  model_name=self.models_names[1], 
+                                  adapter=self.adapter
+                                  scheduler=self.scheduler)
+
+
+        @spaces.GPU
+        def get_first_result(self, img, prompt, negative_prompt, 
+                             controlnet_scale=0.5, strength=1.0,n_steps=30,eta=1.0):
+
+            substrate, resized_image = sketch_process(input_image)
+            prompt=prompt_preprocess(prompt)
+
+            result=self.first_pipe(image=substrate,
+                                    control_image=resized_image,
+                                    strength=strength,
+                                    prompt=prompt,
+                                    negative_prompt = negative_prompt,
+                                    controlnet_conditioning_scale=float(controlnet_scale),
+                                    generator=torch.manual_seed(0),
+                                    num_inference_steps=n_steps,
+                                    eta=eta)
+
+            return result.images[0]
+
+
+        @spaces.GPU
+        def get_second_result(self, img, prompt, negative_prompt,
+                              g_scale=7.5, n_steps=25, 
+                              adapter_scale=0.9, adapter_factor=1.0):
+
+            preprocessed_img=self.detector(img, 
+                                             detect_resolution=1024, 
+                                             image_resolution=1024,
+                                             apply_filter=True).convert("L")
+
+            result=self.second_pipe(prompt=prompt,
+                                    negative_prompt=negative_prompt,
+                                    image=image_preprocessed,
+                                    guidance_scale=g_scale,
+                                    num_inference_steps=n_steps,
+                                    adapter_conditioning_scale=adapter_scale,
+                                    adapter_conditioning_factor=adapter_factor,
+                                    generator = torch.manual_seed(42))
+
+            return result.images[0]
+                
+
+

sd/utils/utils.py

@@ -0,0 +1,78 @@
+import torch
+from diffusers import (ControlNetModel, 
+                       StableDiffusionXLControlNetImg2ImgPipeline, 
+                       AutoencoderKL,
+                       T2IAdapter,
+                       StableDiffusionXLAdapterPipeline,
+                       EulerAncestralDiscreteScheduler)
+
+from controlnet_aux.pidi import PidiNetDetector
+
+from PIL import Image
+import os
+
+
+#VAE=AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+
+#CONTROLNET = ControlNetModel.from_pretrained("diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16)
+
+#ADAPTER = T2IAdapter.from_pretrained("Adapter/t2iadapter", 
+                                    #subfolder="sketch_sdxl_1.0", 
+                                    #torch_dtype=torch.float16, 
+                                    #adapter_type="full_adapter_xl")
+
+
+def get_vae(model_name="madebyollin/sdxl-vae-fp16-fix"):
+    return AutoencoderKL.from_pretrained(model_name, torch_dtype=torch.float16)
+
+def get_controlnet(model_name="diffusers/controlnet-canny-sdxl-1.0"):
+    return ControlNetModel.from_pretrained(model_name, torch_dtype=torch.float16)
+
+def get_adapter(model_name="Adapter/t2iadapter", subfolder="sketch_sdxl_1.0",
+                adapter_type="full_adapter_xl"):
+    if adapter_type == "full_adapter_xl":
+        return T2IAdapter.from_pretrained(model_name, 
+                                    subfolder=subfolder, 
+                                    torch_dtype=torch.float16, 
+                                    adapter_type=adapter_type)
+
+def get_scheduler(model_name, scheduler_type="discrete"):
+    if scheduler_type == "discrete":
+        return EulerAncestralDiscreteScheduler.from_pretrained(model_name, subfolder="scheduler")
+
+
+def get_detector(model_name="lllyasviel/Annotators", model_type='pidi'):
+    if model_type == 'pidi':
+        return PidiNetDetector.from_pretrained(model_name)
+
+
+def load_lora(pipe, lora_path=None):
+    if lora_path != None:
+        try:
+            lora_dir='./'+'/'.join(lora_path.split("/")[:-1])
+            lora_name=lora_path.split("/")[-1]
+            pipe.load_lora_weights(lora_dir, weight_name=lora_name)
+        except Exception as ex:
+            print(ex)
+    #return pipe
+
+
+def get_pipe(vae, model_name, controlnet=None, adapter=None, scheduler=None, lora_path=None):
+    if controlnet!=None:
+        pipe=StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(model_name,
+                                                                          controlnet=controlnet, 
+                                                                          vae=vae, 
+                                                                          torch_dtype=torch.float16)
+
+        load_lora(pipe, lora_path)                                
+        return pipe
+
+    elif adapter != None:
+        pipe=StableDiffusionXLAdapterPipeline.from_pretrained(model_name, 
+                                                                adapter=adapter, 
+                                                                vae=vae,
+                                                                scheduler=scheduler,
+                                                                torch_dtype=torch.float16, 
+                                                                variant="fp16")
+        load_lora(pipe, lora_path)
+        return pipe

ui/gradio_ui.py

@@ -0,0 +1,30 @@
+import gradio as gr
+
+def ui(controller):
+    with gr.Blocks() as ui:
+        with gr.Row():
+            with gr.Column():
+                sketch=gr.Image(sources = 'upload', label='Model image', type = 'pil')
+                first_prompt = gr.Textbox(label="Prompt", lines=3)
+                first_negative_prompt = gr.Textbox(label="Negative prompt", lines=3, value="sketch, lowres, error, extra digit, fewer digits, cropped, worst quality,low quality, normal quality, jpeg artifacts, blurry")
+                #controlnet_scale = gr.Slider(minimum=0.5, maximum=1.25, value=0.5, step=0.01, label="Contr")
+                improve_sketch = gr.Button(value="Improve Sketch", variant="primary")
+            with gr.Column():
+                improved_sketch_view = gr.Image(type="pil", label="Improved Sketch")
+
+        improve_sketch.click(fn=controller.get_first_result,
+                             inputs=[sketch, first_prompt, first_negative_prompt],
+                             outputs=improved_sketch_view)
+
+        with gr.Row():
+            result=gr.Image(type="pil", label="Improved Sketch")
+            second_prompt = gr.Textbox(label="Prompt", lines=3)
+            second_negative_prompt = gr.Textbox(label="Negative prompt", lines=3, value="disfigured, extra digit, fewer digits, cropped, worst quality, low quality")
+            result_button = gr.Button(value="Paint It", variant="primary")
+
+        result_button.click(fn=controller.get_secnd_result,
+                             inputs=[sketch, second_prompt, second_negative_prompt],
+                             outputs=result)
+
+
+    ui.queue().launch(debug=True)

utils/utils.py

@@ -0,0 +1,77 @@
+from PIL import Image
+import numpy as np
+import cv2
+
+
+#first stage sketch preprocess
+def conventional_resize(img):
+    original_width, original_height = img.size    
+    aspect_ratio = original_width / original_height
+
+    conventional_sizes = {
+        1: (1024, 1024),  
+        4/3: (1152, 896), 
+        3/2: (1216, 832),
+        16/9: (1344, 768),
+        21/9: (1568, 672),
+        3/1: (1728, 576),
+        1/4: (512, 2048), 
+        1/3: (576, 1728),
+        9/16: (768, 1344),
+        2/3: (832, 1216),
+        3/4: (896, 1152)
+    }
+
+    closest_aspect_ratio = min(conventional_sizes.keys(), key=lambda x: abs(x - aspect_ratio))
+    new_width, new_height = conventional_sizes[closest_aspect_ratio]
+    
+    resized_img = img.resize((new_width, new_height), Image.LANCZOS)
+
+    return resized_img
+
+
+def get_substrate(img, color=(255, 255, 255, 255)):
+    size=img.size
+    substrate = Image.new("RGBA", size, color) 
+    return substrate.convert("RGB")
+
+
+def sketch_process(img):
+    substrate=conventional_resize(get_substrate(img))
+    resized_img = conventional_resize(img) 
+    return substrate, resized_img
+
+
+
+
+#first stage prompt preprocess
+def remove_duplicates(base_prompt):    
+    prompt_list = base_prompt.split(", ")
+    seen = set()
+    unique_tags = []
+    for tag in prompt_list :
+        tag_clean = tag.lower().strip()
+        if tag_clean not in seen and tag_clean != "":
+            unique_tags.append(tag)
+            seen.add(tag_clean)
+    return ", ".join(unique_tags)
+
+
+def remove_color(base_prompt):    
+    prompt_list = base_prompt.split(", ")
+    color_list = ["pink", "red", "orange", "brown", "yellow", "green", "blue", "purple", "blonde", "colored skin", "white hair"]
+    cleaned_tags = [tag for tag in prompt_list if all(color.lower() not in tag.lower() for color in color_list)]
+    return ", ".join(cleaned_tags)
+
+
+def execute_prompt(base_prompt):
+    prompt_list = base_prompt.split(", ")
+    execute_tags = ["sketch", "transparent background"]    
+    filtered_tags = [tag for tag in prompt_list if tag not in execute_tags]    
+    return ", ".join(filtered_tags)
+
+def prompt_preprocess(prompt):
+    result=execute_prompt(prompt)
+    result=remove_duplicates(result)
+    result=remove_color(result)
+    return result