Spaces:
Paused
Paused
| import os | |
| import re | |
| import math | |
| import tempfile | |
| import zipfile | |
| from pathlib import Path | |
| from typing import List | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from PIL import Image, ImageOps, ImageFilter | |
| from diffusers import ( | |
| ControlNetModel, | |
| StableDiffusionControlNetImg2ImgPipeline, | |
| UniPCMultistepScheduler, | |
| ) | |
| BASE_MODEL = os.getenv("BASE_MODEL", "runwayml/stable-diffusion-v1-5") | |
| CONTROLNET_MODEL = os.getenv("CONTROLNET_MODEL", "lllyasviel/control_v11p_sd15_canny") | |
| DEFAULT_PROMPT = ( | |
| "literal hand-drawn graphite animation frame on warm off-white sketch paper, " | |
| "rough 2B and 4B pencil, visible construction marks, uneven pencil pressure, " | |
| "cross-hatching, eraser ghosts, paper tooth, imperfect hand-drawn contours, " | |
| "rough graphite anime storyboard, preserve the exact composition, pose, camera, " | |
| "character identity and guitar geometry, preserve existing red accents only, " | |
| "no new colors, no digital painting, no smooth vector edges" | |
| ) | |
| DEFAULT_NEGATIVE = ( | |
| "photorealistic, glossy, smooth, digital painting, watercolor, oil paint, airbrush, " | |
| "3d render, vector art, clean ink, extra fingers, extra limbs, deformed hands, " | |
| "mutated anatomy, different scene, different outfit, different guitar, text, watermark, logo" | |
| ) | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 | |
| _pipe = None | |
| def natural_key(text: str): | |
| return [int(c) if c.isdigit() else c.lower() for c in re.split(r"(\d+)", text)] | |
| def load_pipeline(): | |
| global _pipe | |
| if _pipe is not None: | |
| return _pipe | |
| controlnet = ControlNetModel.from_pretrained( | |
| CONTROLNET_MODEL, | |
| torch_dtype=DTYPE, | |
| ) | |
| pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( | |
| BASE_MODEL, | |
| controlnet=controlnet, | |
| torch_dtype=DTYPE, | |
| safety_checker=None, | |
| requires_safety_checker=False, | |
| ) | |
| pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) | |
| pipe.enable_attention_slicing() | |
| if DEVICE == "cuda": | |
| try: | |
| pipe.enable_model_cpu_offload() | |
| except Exception: | |
| pipe = pipe.to("cuda") | |
| else: | |
| pipe = pipe.to("cpu") | |
| _pipe = pipe | |
| return _pipe | |
| def resize_multiple_of_8(img: Image.Image, max_side: int): | |
| img = ImageOps.exif_transpose(img).convert("RGB") | |
| w, h = img.size | |
| scale = min(1.0, float(max_side) / float(max(w, h))) | |
| nw = max(64, int(round((w * scale) / 8) * 8)) | |
| nh = max(64, int(round((h * scale) / 8) * 8)) | |
| return img.resize((nw, nh), Image.Resampling.LANCZOS) | |
| def canny_control(img: Image.Image): | |
| # Dependency-light edge map: PIL + NumPy only. | |
| gray = ImageOps.grayscale(img) | |
| edges = gray.filter(ImageFilter.FIND_EDGES) | |
| arr = np.asarray(edges, dtype=np.uint8) | |
| # suppress weak paper/grain edges so composition dominates | |
| arr = np.where(arr > 28, 255, 0).astype(np.uint8) | |
| edge = Image.fromarray(arr, mode="L") | |
| return Image.merge("RGB", (edge, edge, edge)) | |
| def extract_images(zip_path: str, workdir: str) -> List[Path]: | |
| src = Path(workdir) / "input" | |
| src.mkdir(parents=True, exist_ok=True) | |
| with zipfile.ZipFile(zip_path, "r") as zf: | |
| zf.extractall(src) | |
| exts = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} | |
| imgs = [p for p in src.rglob("*") if p.suffix.lower() in exts] | |
| imgs.sort(key=lambda p: natural_key(str(p.relative_to(src)))) | |
| if not imgs: | |
| raise gr.Error("No PNG/JPG/WebP/BMP images found inside the ZIP.") | |
| return imgs | |
| def make_sheet(paths, out_path: Path, cols=4, thumb=(240, 180)): | |
| if not paths: | |
| return None | |
| canvases = [] | |
| for p in paths[:16]: | |
| im = Image.open(p).convert("RGB") | |
| im.thumbnail(thumb, Image.Resampling.LANCZOS) | |
| bg = Image.new("RGB", thumb, (246, 243, 236)) | |
| bg.paste(im, ((thumb[0]-im.width)//2, (thumb[1]-im.height)//2)) | |
| canvases.append(bg) | |
| rows = math.ceil(len(canvases) / cols) | |
| sheet = Image.new("RGB", (cols*thumb[0], rows*thumb[1]), (250,248,242)) | |
| for i, im in enumerate(canvases): | |
| sheet.paste(im, ((i % cols)*thumb[0], (i // cols)*thumb[1])) | |
| sheet.save(out_path) | |
| return str(out_path) | |
| def process_zip( | |
| zip_file, | |
| prompt, | |
| negative_prompt, | |
| max_side, | |
| strength, | |
| control_scale, | |
| steps, | |
| guidance, | |
| seed, | |
| progress=gr.Progress(), | |
| ): | |
| if zip_file is None: | |
| raise gr.Error("Upload a ZIP first.") | |
| pipe = load_pipeline() | |
| workdir = tempfile.mkdtemp(prefix="graphite_") | |
| out_dir = Path(workdir) / "output" | |
| out_dir.mkdir(parents=True, exist_ok=True) | |
| images = extract_images(zip_file, workdir) | |
| previews = [] | |
| for idx, path in enumerate(images, start=1): | |
| progress((idx - 1) / len(images), desc=f"{idx}/{len(images)} · {path.name}") | |
| source = resize_multiple_of_8(Image.open(path), int(max_side)) | |
| control = canny_control(source) | |
| generator = torch.Generator(device="cpu").manual_seed(int(seed) + idx) | |
| result = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| image=source, | |
| control_image=control, | |
| strength=float(strength), | |
| controlnet_conditioning_scale=float(control_scale), | |
| num_inference_steps=int(steps), | |
| guidance_scale=float(guidance), | |
| generator=generator, | |
| ).images[0] | |
| out = out_dir / f"{path.stem}.png" | |
| result.save(out, format="PNG") | |
| if len(previews) < 16: | |
| previews.append(out) | |
| result_zip = Path(workdir) / "graphite_frames.zip" | |
| with zipfile.ZipFile(result_zip, "w", zipfile.ZIP_DEFLATED) as zf: | |
| for p in sorted(out_dir.glob("*.png"), key=lambda x: natural_key(x.name)): | |
| zf.write(p, arcname=p.name) | |
| sheet = make_sheet(previews, Path(workdir) / "preview.png") | |
| progress(1.0, desc="Done") | |
| status = ( | |
| f"Processed {len(images)} frame(s).\n" | |
| f"Runtime device: {DEVICE}\n" | |
| f"Base model: {BASE_MODEL}\n" | |
| f"ControlNet: {CONTROLNET_MODEL}" | |
| ) | |
| return str(result_zip), sheet, status | |
| with gr.Blocks(title="Graphite Frame Batch Restorer") as demo: | |
| gr.Markdown( | |
| "# ✏️ Graphite Frame Batch Restorer\n" | |
| "Upload a ZIP of individual storyboard/flip-book frames. " | |
| "The app redraws each frame as rough graphite while using an edge-control pass " | |
| "to keep composition and pose from wandering." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| zip_in = gr.File(label="ZIP of frames", file_types=[".zip"], type="filepath") | |
| prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=6) | |
| negative = gr.Textbox(label="Negative prompt", value=DEFAULT_NEGATIVE, lines=4) | |
| with gr.Accordion("Advanced", open=False): | |
| max_side = gr.Slider(512, 1024, value=768, step=64, label="Model max side") | |
| strength = gr.Slider(0.15, 0.60, value=0.28, step=0.01, label="Img2Img strength") | |
| control_scale = gr.Slider(0.50, 1.50, value=1.00, step=0.05, label="Structure lock") | |
| steps = gr.Slider(10, 35, value=22, step=1, label="Steps") | |
| guidance = gr.Slider(3.0, 10.0, value=6.5, step=0.5, label="Guidance") | |
| seed = gr.Number(value=1337, precision=0, label="Seed") | |
| run = gr.Button("Process ZIP", variant="primary") | |
| with gr.Column(): | |
| zip_out = gr.File(label="Processed ZIP") | |
| preview = gr.Image(label="Preview") | |
| status = gr.Textbox(label="Status", lines=5) | |
| run.click( | |
| process_zip, | |
| inputs=[zip_in, prompt, negative, max_side, strength, control_scale, steps, guidance, seed], | |
| outputs=[zip_out, preview, status], | |
| ) | |
| demo.queue(default_concurrency_limit=1).launch(ssr_mode=False) | |