Model Card for UltraSketch

UltraSketch is a diffusion model that has been trained primarily to convert scientific figures into sketches with a hand-drawn style. It is based on SD3_UltraEdit_w_mask and was fine-tuned on SketchFig. In addition, figures from DaTikZ_v2 rendered in a hand-drawn style using Rough.js, as well as the Sketchy Database and Photo Sketching datasets, have been used for data augmentation. Check out the DeTikZify project for more information.

Usage

from PIL import Image
from datasets import load_dataset
from diffusers import DiffusionPipeline
import torch

figure = load_dataset("nllg/datikz-v2", split="train")['image'][0]

pipe = DiffusionPipeline.from_pretrained(
    pretrained_model_name_or_path="nllg/ultrasketch",
    custom_pipeline="nllg/ultrasketch",
    trust_remote_code=True,
    torch_dtype=torch.float16,
    device_map="balanced"
)

sketch = pipe(
    prompt="Turn it into a hand-drawn sketch",
    image=figure,
    mask_img=Image.new("RGB", figure.size, "white"),
    num_inference_steps=50,
    image_guidance_scale=1.7,
    guidance_scale=1.5,
    strength=0.9
).images[0]

sketch.save("sketch.png")