Spaces:
Sleeping
Sleeping
File size: 2,623 Bytes
5086d03 05a2de2 5086d03 8c9e7fc 5086d03 146ebdb 5086d03 146ebdb 5086d03 146ebdb 5086d03 146ebdb 5086d03 cdc3dcf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import hashlib
import io
import torch
from pathlib import Path
from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, UniPCMultistepScheduler
from PIL import Image, ImageOps
import gradio as gr
# ---- Model loading ----
CACHE_DIR = "./cache"
CNET_MODEL = "MrPio/Texture-Anything_CNet-SD15"
SD_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
controlnet = ControlNetModel.from_pretrained(
CNET_MODEL, cache_dir=CACHE_DIR, torch_dtype=torch.float16
)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
SD_MODEL,
controlnet=controlnet,
cache_dir=CACHE_DIR,
torch_dtype=torch.float16,
safety_checker=None,
)
# speed & memory optimizations
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
# pipe.enable_xformers_memory_efficient_attention() # if xformers installed
# pipe.enable_model_cpu_offload()
def pil2hash(image: Image.Image) -> str:
buffer = io.BytesIO()
image.save(buffer, format="PNG")
image_bytes = buffer.getvalue()
return hashlib.sha256(image_bytes).hexdigest()
def caption2hash(caption: str) -> str:
return hashlib.sha256(caption.encode()).hexdigest()
# ---- Inference function ----
def infer(caption: str, condition_image: Image.Image, steps: int = 20, seed: int = 0, invert: bool = False):
print("Loading condition image")
img = condition_image.convert("RGB")
if invert:
img = ImageOps.invert(img)
print("Condition image inverted")
cache_file = Path(f"inferences/{pil2hash(img)}_{caption2hash(caption)}.png")
if cache_file.exists():
return Image.open(cache_file)
generator = torch.manual_seed(seed)
print("Starting generation...")
output = pipe(prompt=caption, image=img, num_inference_steps=steps, generator=generator).images[0]
print("Caching result...")
output.save(cache_file)
return output
# ---- Gradio UI + API ----
with gr.Blocks() as demo:
gr.Markdown("## ControlNet + Stable Diffusion 1.5")
with gr.Row():
txt = gr.Textbox(label="Prompt", placeholder="Describe the texture...")
cond = gr.Image(type="pil", label="Condition Image")
with gr.Row():
steps = gr.Slider(1, 50, value=20, label="Inference Steps")
seed = gr.Number(value=0, label="Seed (0 for random)")
inv = gr.Checkbox(label="Invert UV colors?")
btn = gr.Button("Generate")
out = gr.Image(label="Output")
btn.click(fn=infer, inputs=[txt, cond, steps, seed, inv], outputs=out)
# enable the standard gradio REST API (/run/predict)
demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
|