Initial commit

app.py

@@ -0,0 +1,55 @@
+import os
+
+import gradio as gr
+from inference import load_models, cache_path
+from PIL import Image
+from os import path
+
+canvas_size = 512
+
+if not path.exists(cache_path):
+    os.makedirs(cache_path, exist_ok=True)
+
+with gr.Blocks() as demo:
+    infer = load_models()
+
+    with gr.Column():
+        with gr.Row():
+            with gr.Column():
+                s = gr.Slider(label="steps", minimum=4, maximum=8, step=1, value=4, interactive=True)
+                c = gr.Slider(label="cfg", minimum=0.1, maximum=3, step=0.1, value=1, interactive=True)
+                i_s = gr.Slider(label="sketch strength", minimum=0.1, maximum=0.9, step=0.1, value=0.9, interactive=True)
+            with gr.Column():
+                mod = gr.Text(label="Model HuggingFace id (after changing this wait until the model downloads in the console)", value="Lykon/dreamshaper-7", interactive=True)
+                t = gr.Text(label="Prompt", value="Scary warewolf, 8K, realistic, colorful, long sharp teeth, splash art", interactive=True)
+                se = gr.Number(label="seed", value=1337, interactive=True)
+        with gr.Row(equal_height=True):
+            i = gr.Image(source="canvas", tool="color-sketch", shape=(canvas_size, canvas_size), width=canvas_size, height=canvas_size, type="pil")
+            o = gr.Image(width=canvas_size, height=canvas_size)
+
+            def process_image(p, im, steps, cfg, image_strength, seed):
+                if not im:
+                    return Image.new("RGB", (canvas_size, canvas_size))
+                return infer(
+                    prompt=p,
+                    image=im,
+                    num_inference_steps=steps,
+                    guidance_scale=cfg,
+                    strength=image_strength,
+                    seed=int(seed)
+                )
+
+            reactive_controls = [t, i, s, c, i_s, se]
+
+            for control in reactive_controls:
+                control.change(fn=process_image, inputs=reactive_controls, outputs=o)
+
+            def update_model(model_name):
+                global infer
+                infer = load_models(model_name)
+
+            mod.change(fn=update_model, inputs=mod)
+
+
+if __name__ == "__main__":
+    demo.launch()

inference.py

@@ -0,0 +1,103 @@
+import os
+import random
+from os import path
+from contextlib import nullcontext
+import time
+from sys import platform
+import torch
+
+cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
+
+os.environ["TRANSFORMERS_CACHE"] = cache_path
+os.environ["HF_HUB_CACHE"] = cache_path
+os.environ["HF_HOME"] = cache_path
+is_mac = platform == "darwin"
+
+def should_use_fp16():
+    if is_mac:
+        return True
+
+    gpu_props = torch.cuda.get_device_properties("cuda")
+
+    if gpu_props.major < 6:
+        return False
+
+    nvidia_16_series = ["1660", "1650", "1630"]
+
+    for x in nvidia_16_series:
+        if x in gpu_props.name:
+            return False
+
+    return True
+
+class timer:
+    def __init__(self, method_name="timed process"):
+        self.method = method_name
+
+    def __enter__(self):
+        self.start = time.time()
+        print(f"{self.method} starts")
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        end = time.time()
+        print(f"{self.method} took {str(round(end - self.start, 2))}s")
+
+
+def load_models(model_id="Lykon/dreamshaper-7"):
+    from diffusers import AutoPipelineForImage2Image, LCMScheduler
+    from diffusers.utils import load_image
+
+    if not is_mac:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    use_fp16 = should_use_fp16()
+
+    lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"
+
+    if use_fp16:
+        pipe = AutoPipelineForImage2Image.from_pretrained(
+            model_id,
+            cache_dir=cache_path,
+            torch_dtype=torch.float16,
+            variant="fp16",
+            safety_checker=None
+        )
+    else:
+        pipe = AutoPipelineForImage2Image.from_pretrained(
+            model_id,
+            cache_dir=cache_path,
+            safety_checker=None
+        )
+
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+
+    pipe.load_lora_weights(lcm_lora_id)
+    pipe.fuse_lora()
+
+    device = "mps" if is_mac else "cuda"
+
+    pipe.to(device=device)
+
+    generator = torch.Generator()
+
+    def infer(
+            prompt,
+            image,
+            num_inference_steps=4,
+            guidance_scale=1,
+            strength=0.9,
+            seed=random.randrange(0, 2**63)
+    ):
+        with torch.inference_mode():
+            with torch.autocast("cuda") if device == "cuda" else nullcontext():
+                with timer("inference"):
+                    return pipe(
+                        prompt=prompt,
+                        image=load_image(image),
+                        generator=generator.manual_seed(seed),
+                        num_inference_steps=num_inference_steps,
+                        guidance_scale=guidance_scale,
+                        strength=strength
+                    ).images[0]
+
+    return infer

requirements.txt

@@ -0,0 +1,7 @@
+diffusers
+flask
+torch
+transformers
+accelerate
+pillow
+gradio==3.41.2