ultrasketch / README.md
potamides's picture
Update README.md
000df51
metadata
library_name: diffusers
license: other
base_model: BleachNick/SD3_UltraEdit_w_mask

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 DaTikZv2 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")