Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces | |
| import pickle | |
| import numpy as np | |
| import torch | |
| import gradio as gr | |
| import imageio.v2 as imageio | |
| from PIL import Image | |
| from huggingface_hub import hf_hub_download | |
| from model import AlignDRAW | |
| REPO = "mediasynthesismuseum/aligndraw" | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| OUTPUT_SIZE = 512 | |
| # Original trained architecture and weights | |
| model = AlignDRAW(imgSize=32, runSteps=32) | |
| checkpoint_path = hf_hub_download( | |
| REPO, | |
| "aligndraw_coco.pt" | |
| ) | |
| checkpoint = torch.load( | |
| checkpoint_path, | |
| map_location="cpu" | |
| ) | |
| model.load_state_dict(checkpoint) | |
| model.eval().to("cuda") | |
| with open( | |
| hf_hub_download(REPO, "dictionary.pkl"), | |
| "rb" | |
| ) as f: | |
| DICT = pickle.load(f, encoding="latin1") | |
| def tokenize(caption): | |
| words = caption.strip().lower().replace(".", " .").split() | |
| unknown = sorted({ | |
| word for word in words | |
| if word not in DICT and word != "." | |
| }) | |
| tokens = [ | |
| DICT.get(word, DICT["UNK"]) | |
| for word in words | |
| ] or [DICT["UNK"]] | |
| return tokens, unknown | |
| def upscale(tensor): | |
| array = ( | |
| tensor.detach() | |
| .float() | |
| .cpu() | |
| .permute(1, 2, 0) | |
| .numpy() * 255 | |
| ).clip(0, 255).astype(np.uint8) | |
| image = Image.fromarray(array, mode="RGB") | |
| image = image.resize( | |
| (OUTPUT_SIZE, OUTPUT_SIZE), | |
| Image.Resampling.NEAREST | |
| ) | |
| assert image.size == (512, 512) | |
| return image | |
| def generate(caption, n_samples): | |
| if not caption or not caption.strip(): | |
| raise gr.Error("Type a caption first.") | |
| tokens, unknown = tokenize(caption) | |
| y = torch.tensor( | |
| [tokens], | |
| dtype=torch.long, | |
| device="cuda" | |
| ) | |
| with torch.inference_mode(): | |
| images = model.generate( | |
| y.expand(int(n_samples), -1) | |
| ) | |
| gallery = [ | |
| upscale(images[i]) | |
| for i in range(int(n_samples)) | |
| ] | |
| _, sequence = model.generate( | |
| y, | |
| return_seq=True | |
| ) | |
| frames = [ | |
| np.array(upscale(step[0])) | |
| for step in sequence | |
| ] | |
| # Keep original 32 drawing frames, then hold the final result | |
| if frames: | |
| while len(frames) < 69: | |
| frames.append(frames[-1]) | |
| frames.extend([frames[-1]] * 8) | |
| gif_path = os.path.join(HERE, "draw.gif") | |
| imageio.mimsave( | |
| gif_path, | |
| frames, | |
| duration=0.09, | |
| loop=0 | |
| ) | |
| note = ( | |
| "💡 Unknown vocabulary words (treated as UNK): " | |
| + ", ".join(unknown) | |
| ) if unknown else "" | |
| print("OUTPUT DIMENSIONS:", gallery[0].size) | |
| return gallery, gif_path, note | |
| EXAMPLES = [ | |
| "a red stop sign on a street corner .", | |
| "a herd of elephants walking in a field .", | |
| "a yellow school bus on the road .", | |
| "a bowl of oranges on a table .", | |
| "a plate with a sandwich and french fries .", | |
| "a large blue passenger jet in the sky .", | |
| ] | |
| CSS = """ | |
| .gradio-container { | |
| max-width: 1100px !important; | |
| margin: 0 auto !important; | |
| } | |
| #ad-title { | |
| text-align: center; | |
| } | |
| """ | |
| with gr.Blocks(title="AlignDRAW 512") as demo: | |
| gr.Markdown( | |
| "# 🖍️ AlignDRAW — 512×512 Output\n" | |
| "The original COCO-trained AlignDRAW model, " | |
| "with images enlarged to exactly 512×512 pixels.\n\n" | |
| "Original checkpoint and vocabulary preserved.", | |
| elem_id="ad-title" | |
| ) | |
| with gr.Row(): | |
| caption = gr.Textbox( | |
| label="Caption", | |
| value=EXAMPLES[0], | |
| scale=4 | |
| ) | |
| samples = gr.Slider( | |
| 1, 8, 4, | |
| step=1, | |
| label="Samples" | |
| ) | |
| draw = gr.Button("🖍️ Draw", variant="primary") | |
| note = gr.Markdown() | |
| gallery = gr.Gallery( | |
| label="Generated images — 512×512", | |
| columns=2, | |
| height=550 | |
| ) | |
| animation = gr.Image( | |
| label="AlignDRAW drawing process", | |
| height=350 | |
| ) | |
| draw.click( | |
| generate, | |
| inputs=[caption, samples], | |
| outputs=[gallery, animation, note] | |
| ) | |
| gr.Examples( | |
| EXAMPLES, | |
| inputs=caption | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue().launch( | |
| theme=gr.themes.Citrus(), | |
| css=CSS | |
| ) | |