Initial commit

.gitignore

@@ -0,0 +1,10 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.python-version

@@ -0,0 +1 @@
+3.11

README.md

@@ -1,14 +0,0 @@
----
-title: ColorMasking
-emoji: 🌖
-colorFrom: purple
-colorTo: pink
-sdk: gradio
-sdk_version: 5.14.0
-app_file: app.py
-pinned: false
-license: mit
-short_description: A very simple demonstration of color masking & diffusion
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,9 @@
+from models import LineartGenerator
+from ui import SketchToImageApp
+
+
+if __name__ == "__main__":
+    lineart_generator = LineartGenerator()
+    app = SketchToImageApp(lineart_generator).create_interface()
+    app.launch()
+

models.py

@@ -0,0 +1,80 @@
+import numpy as np
+import torch
+from PIL import Image
+from controlnet_aux import LineartDetector
+from diffusers import (
+    ControlNetModel,
+    StableDiffusionControlNetPipeline,
+    UniPCMultistepScheduler
+)
+
+
+class LineartGenerator:
+
+    def __init__(self, device: str = None, seed: int = 0):
+        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
+        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
+        self.seed = seed
+        self._initialize_models()
+
+    def _initialize_models(self):
+        self.lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
+        checkpoint = "ControlNet-1-1-preview/control_v11p_sd15_lineart"
+        self.controlnet = ControlNetModel.from_pretrained(checkpoint)
+        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5",
+            controlnet=self.controlnet
+        )
+        self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
+        self.pipe.to(self.device)
+
+    @staticmethod
+    def load_lineart_image(uploaded_file, size: tuple = (512, 512)) -> np.ndarray:
+        if not uploaded_file:
+            return None
+
+        with open(uploaded_file.name, "rb") as file_obj:
+            image = Image.open(file_obj).convert("L")
+            image = image.resize(size)
+        return np.array(image, dtype=np.uint8)
+
+    @staticmethod
+    def merge_lineart_and_brush(brush_canvas, uploaded_file, size: tuple = (512, 512)) -> Image.Image:
+        if brush_canvas is None or uploaded_file is None:
+            return None
+
+        # Reload and process the original lineart image
+        with open(uploaded_file.name, "rb") as file_obj:
+            lineart_image = Image.open(file_obj).convert("L")
+            lineart_image = lineart_image.resize(size)
+
+        lineart_rgba = lineart_image.convert("RGBA")
+        processed_pixels = []
+        for pixel in lineart_rgba.getdata():
+            if pixel[0] > 240:
+                processed_pixels.append((255, 255, 255, 0))
+            else:
+                processed_pixels.append((0, 0, 0, 255))
+        lineart_rgba.putdata(processed_pixels)
+
+        brush_layer = Image.fromarray(brush_canvas["composite"]).convert("RGBA")
+
+        combined = Image.alpha_composite(brush_layer, lineart_rgba)
+        return combined
+
+    def generate_image(self, annotated_lineart: Image.Image,
+                       prompt: str = "",
+                       num_inference_steps: int = 30) -> Image.Image:
+        if annotated_lineart is None:
+            raise ValueError("No annotated lineart provided!")
+
+        annotated_lineart = annotated_lineart.resize((512, 512))
+        refined_lineart = self.lineart_detector(annotated_lineart).convert("RGBA")
+        generator = torch.manual_seed(self.seed)
+
+        return self.pipe(
+            prompt=prompt,
+            image=refined_lineart,
+            num_inference_steps=num_inference_steps,
+            generator=generator
+        ).images[0]

pyproject.toml

@@ -0,0 +1,7 @@
+[project]
+name = "color-mapping"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = []

requirements.txt

@@ -0,0 +1,221 @@
+aif360==0.6.1
+aiohttp==3.9.5
+aiosignal==1.3.1
+altair==5.3.0
+annotated-types==0.7.0
+anyio==4.3.0
+appnope==0.1.4
+argon2-cffi==23.1.0
+argon2-cffi-bindings==21.2.0
+arrow==1.3.0
+association-metrics==0.0.1
+asttokens==2.4.1
+async-lru==2.0.4
+attrs==23.2.0
+Babel==2.15.0
+beautifulsoup4==4.12.3
+bleach==6.1.0
+blinker==1.8.2
+boto3==1.34.135
+botocore==1.34.135
+cachetools==5.3.3
+certifi==2024.2.2
+cffi==1.16.0
+charset-normalizer==3.3.2
+click==8.1.7
+comm==0.2.2
+contourpy==1.2.1
+cssutils==2.11.0
+cycler==0.12.1
+dataclasses-json==0.6.6
+dataframe-image==0.2.3
+debugpy==1.8.1
+decorator==5.1.1
+defusedxml==0.7.1
+diskcache==5.6.3
+dnspython==2.6.1
+email_validator==2.1.1
+executing==2.0.1
+faiss-cpu==1.8.0
+fastapi==0.111.0
+fastapi-cli==0.0.4
+fastjsonschema==2.19.1
+filelock==3.14.0
+fonttools==4.52.4
+fqdn==1.5.1
+frozenlist==1.4.1
+fsspec==2024.5.0
+gensim==4.3.2
+gitdb==4.0.11
+GitPython==3.1.43
+greenlet==3.0.3
+h11==0.14.0
+html2image==2.0.4.3
+httpcore==1.0.5
+httptools==0.6.1
+httpx==0.27.0
+huggingface-hub==0.23.2
+idna==3.7
+imbalanced-learn==0.12.3
+imblearn==0.0
+iniconfig==2.0.0
+ipykernel==6.29.4
+ipython==8.24.0
+ipywidgets==8.1.2
+isoduration==20.11.0
+jedi==0.19.1
+Jinja2==3.1.4
+jmespath==1.0.1
+joblib==1.4.2
+json5==0.9.25
+jsonpatch==1.33
+jsonpointer==2.4
+jsonschema==4.22.0
+jsonschema-specifications==2023.12.1
+jupyter==1.0.0
+jupyter-console==6.6.3
+jupyter-events==0.10.0
+jupyter-lsp==2.2.5
+jupyter_client==8.6.1
+jupyter_core==5.7.2
+jupyter_server==2.14.0
+jupyter_server_terminals==0.5.3
+jupyterlab==4.1.8
+jupyterlab_pygments==0.3.0
+jupyterlab_server==2.27.1
+jupyterlab_widgets==3.0.10
+kiwisolver==1.4.5
+langchain==0.2.1
+langchain-community==0.2.1
+langchain-core==0.2.3
+langchain-text-splitters==0.2.0
+langchainhub==0.1.17
+langsmith==0.1.67
+Levenshtein==0.25.1
+llama_cpp_python==0.2.76
+lxml==5.2.2
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+marshmallow==3.21.2
+matplotlib==3.9.0
+matplotlib-inline==0.1.7
+mdurl==0.1.2
+mistune==3.0.2
+mpmath==1.3.0
+multidict==6.0.5
+mypy-extensions==1.0.0
+nbclient==0.10.0
+nbconvert==7.16.4
+nbformat==5.10.4
+nest-asyncio==1.6.0
+networkx==3.3
+notebook==7.1.3
+notebook_shim==0.2.4
+numpy==1.26.4
+orjson==3.10.3
+overrides==7.7.0
+packaging==23.2
+pandas==2.2.2
+pandocfilters==1.5.1
+parso==0.8.4
+patsy==0.5.6
+pexpect==4.9.0
+pika==1.3.2
+pillow==10.3.0
+platformdirs==4.2.1
+plotly==5.22.0
+pluggy==1.5.0
+polars==0.20.31
+prometheus_client==0.20.0
+prompt-toolkit==3.0.43
+protobuf==4.25.3
+psutil==5.9.8
+psycopg2==2.9.10
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pyaml==24.4.0
+pyarrow==16.1.0
+pycparser==2.22
+pydantic==2.10.3
+pydantic_core==2.27.1
+pydash==8.0.1
+pydeck==0.9.1
+Pygments==2.18.0
+pykalman==0.9.7
+pyparsing==3.1.2
+pytest==8.2.1
+pytest-mock==3.14.0
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-json-logger==2.0.7
+python-Levenshtein==0.25.1
+python-multipart==0.0.9
+pytz==2024.1
+PyWavelets==1.6.0
+PyYAML==6.0.1
+pyzmq==26.0.3
+qtconsole==5.5.2
+QtPy==2.4.1
+rapidfuzz==3.9.0
+redis==5.2.1
+referencing==0.35.1
+regex==2024.5.15
+requests==2.31.0
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rich==13.7.1
+rpds-py==0.18.1
+s3transfer==0.10.2
+safetensors==0.4.3
+scikit-learn==1.4.2
+scikit-optimize==0.10.2
+scipy==1.13.0
+seaborn==0.13.2
+Send2Trash==1.8.3
+sentence-transformers==3.0.0
+shellingham==1.5.4
+simpy==4.1.1
+six==1.16.0
+smart-open==7.0.4
+smmap==5.0.1
+sniffio==1.3.1
+soupsieve==2.5
+SQLAlchemy==2.0.36
+stack-data==0.6.3
+starlette==0.37.2
+statsmodels==0.14.2
+streamlit==1.35.0
+sympy==1.12.1
+tenacity==8.3.0
+terminado==0.18.1
+threadpoolctl==3.5.0
+tinycss2==1.3.0
+tokenizers==0.19.1
+toml==0.10.2
+toolz==0.12.1
+torch==2.2.2
+tornado==6.4
+tqdm==4.66.4
+traitlets==5.14.3
+transformers==4.41.2
+typer==0.12.3
+types-python-dateutil==2.9.0.20240316
+types-requests==2.32.0.20240602
+typing-inspect==0.9.0
+typing_extensions==4.12.2
+tzdata==2024.1
+ujson==5.10.0
+uri-template==1.3.0
+urllib3==2.2.1
+uvicorn==0.30.0
+uvloop==0.19.0
+watchfiles==0.22.0
+wcwidth==0.2.13
+webcolors==1.13
+webencodings==0.5.1
+websocket-client==1.8.0
+websockets==12.0
+widgetsnbextension==4.0.10
+wrapt==1.16.0
+yarl==1.9.4
+controlnet_aux

space.yaml

@@ -0,0 +1,2 @@
+sdk: gradio
+accelerator: gpu

ui.py

@@ -0,0 +1,57 @@
+import gradio as gr
+from PIL import Image
+from models import LineartGenerator
+
+
+class SketchToImageApp:
+    """
+    An application that combines a user’s annotated sketch with original lineart,
+    and then generates an image using a ControlNet-based pipeline.
+    """
+
+    def __init__(self, lineart_generator: LineartGenerator):
+        self.lineart_generator = lineart_generator
+
+    def generate_image(self, brush_canvas: dict, uploaded_file) -> Image.Image:
+        merged_lineart = self.lineart_generator.merge_lineart_and_brush(brush_canvas, uploaded_file)
+        return self.lineart_generator.generate_image(merged_lineart, num_inference_steps=30)
+
+    def create_interface(self) -> gr.Blocks:
+        with gr.Blocks() as app:
+            gr.Markdown(
+                "# Lineart & Color Mask With Controlnet\n"
+                "Brush strokes will be applied behind the processed lineart so that the "
+                "black lines always remain visible."
+            )
+
+            lineart_file_input = gr.File(
+                label="Upload Lineart Sketch",
+                file_types=["image"],
+                file_count="single"
+            )
+
+            with gr.Row():
+                brush_canvas_input = gr.Sketchpad(
+                    label="Annotate Your Lineart",
+                    type="numpy",
+                    brush=gr.Brush(),
+                    width=512,
+                    height=512
+                )
+
+            lineart_file_input.change(
+                fn=self.lineart_generator.load_lineart_image,
+                inputs=lineart_file_input,
+                outputs=brush_canvas_input
+            )
+
+            generate_button = gr.Button("Generate")
+            output_image = gr.Image(label="Generated Image")
+
+            generate_button.click(
+                fn=self.generate_image,
+                inputs=[brush_canvas_input, lineart_file_input],
+                outputs=output_image
+            )
+
+        return app

uv.lock

@@ -0,0 +1,7 @@
+version = 1
+requires-python = ">=3.11"
+
+[[package]]
+name = "color-mapping"
+version = "0.1.0"
+source = { virtual = "." }