make demo

app.py

@@ -1,7 +1,63 @@
+import torch
 import gradio as gr
+from torchvision.utils import make_grid
+from torchvision.transforms.v2.functional import to_pil_image
+from huggingface_hub import hf_hub_download
+from digitdreamer import Autoencoder, DiT
+from digitdreamer.modules import RF
+from PIL.Image import Image
 
-def greet(name):
-    return "Hello " + name + "!!"
+hf_hub_download(
+    "karanravindra/digitdreamer", "ft-decoder.pth", subfolder="models", local_dir="."
+)
+hf_hub_download(
+    "karanravindra/digitdreamer", "diffusion.pth", subfolder="models", local_dir="."
+)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+torch.set_grad_enabled(False)
+decoder = Autoencoder().decoder
+dit = DiT()
+
+decoder.load_state_dict(torch.load("models/ft-decoder.pth", weights_only=True))
+dit.load_state_dict(torch.load("models/diffusion.pth", weights_only=True))
+
+rf = RF(dit)
+
+
+def generate(choice: str, images: int, steps: int, cfg: float):
+    if choice != "Random":
+        class_choice = int(choice) + 1
+        cond = torch.full((images,), class_choice, dtype=torch.long)
+    else:
+        class_choice = torch.randint(1, 11, (images,))
+        cond = class_choice
+
+    noise = torch.randn(images, 8, 2, 2)
+    uncond = torch.full((images,), 0, dtype=torch.long)
+
+    samples = rf.sample(noise, cond, uncond, sample_steps=steps, cfg=cfg)
+
+    samples = torch.cat(samples, dim=0)
+
+    imgs = decoder(samples).cpu()
+    imgs = imgs.view(-1, images, 1, 32, 32)
+
+    pil_imgs: list[Image] = [to_pil_image(make_grid(img, nrow=10)) for img in imgs]
+    
+    return pil_imgs[-1]
+
+
+demo = gr.Interface(
+    fn=generate,
+    submit_btn="Generate",
+    inputs=gr.Radio(label="Number", choices=list("0123456789")+["Random"], value="Random"),
+    additional_inputs=[
+        gr.Slider(label="Number of Images", minimum=10, maximum=100, step=10, value=100),
+        gr.Slider(label="Number of Steps", minimum=1, maximum=100, step=1, value=6),
+        gr.Slider(label="Classifier Free Guidence", minimum=0, maximum=10, step=0.1, value=2)
+    ],
+    outputs=gr.Image(),
+    title="DigitDreamer",
+    description="Generate images of a number using the DiT model",
+)
 demo.launch()

requirements.txt

@@ -0,0 +1,114 @@
+digitdreamer @ git+https://github.com/karanravindra/digitdreamer@main
+huggingface_hub
+
+# This file was autogenerated via `uv export`.
+aiofiles==23.2.1
+annotated-types==0.7.0
+anyio==4.6.2.post1
+appnope==0.1.4 ; platform_system == 'Darwin'
+asttokens==2.4.1
+certifi==2024.8.30
+cffi==1.17.1 ; implementation_name == 'pypy'
+charset-normalizer==3.4.0
+click==8.1.7 ; sys_platform != 'emscripten'
+colorama==0.4.6 ; sys_platform == 'win32' or platform_system == 'Windows'
+comm==0.2.2
+contourpy==1.3.0
+cycler==0.12.1
+debugpy==1.8.7
+decorator==5.1.1
+einops==0.8.0
+exceptiongroup==1.2.2 ; python_full_version < '3.11'
+executing==2.1.0
+fastapi==0.115.4
+ffmpy==0.4.0
+filelock==3.16.1
+fonttools==4.54.1
+fsspec==2024.10.0
+gradio==5.3.0
+gradio-client==1.4.2
+h11==0.14.0
+httpcore==1.0.6
+httpx==0.27.2
+huggingface-hub==0.26.2
+idna==3.10
+ipykernel==6.29.5
+ipython==8.29.0
+ipywidgets==8.1.5
+jedi==0.19.1
+jinja2==3.1.4
+jupyter-client==8.6.3
+jupyter-core==5.7.2
+jupyterlab-widgets==3.0.13
+kiwisolver==1.4.7
+markdown-it-py==3.0.0 ; sys_platform != 'emscripten'
+markupsafe==2.1.5
+matplotlib==3.9.2
+matplotlib-inline==0.1.7
+mdurl==0.1.2 ; sys_platform != 'emscripten'
+mpmath==1.3.0
+nest-asyncio==1.6.0
+networkx==3.4.2
+numpy==2.1.2
+nvidia-cublas-cu12==12.4.5.8 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cuda-cupti-cu12==12.4.127 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cuda-nvrtc-cu12==12.4.127 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cuda-runtime-cu12==12.4.127 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cudnn-cu12==9.1.0.70 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cufft-cu12==11.2.1.3 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-curand-cu12==10.3.5.147 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cusolver-cu12==11.6.1.9 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cusparse-cu12==12.3.1.170 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-nccl-cu12==2.21.5 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-nvjitlink-cu12==12.4.127 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-nvtx-cu12==12.4.127 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+orjson==3.10.10
+packaging==24.1
+pandas==2.2.3
+parso==0.8.4
+pexpect==4.9.0 ; sys_platform != 'emscripten' and sys_platform != 'win32'
+pillow==10.4.0
+platformdirs==4.3.6
+prompt-toolkit==3.0.48
+psutil==6.1.0
+ptyprocess==0.7.0 ; sys_platform != 'emscripten' and sys_platform != 'win32'
+pure-eval==0.2.3
+pycparser==2.22 ; implementation_name == 'pypy'
+pydantic==2.9.2
+pydantic-core==2.23.4
+pydub==0.25.1
+pygments==2.18.0
+pyparsing==3.2.0
+python-dateutil==2.9.0.post0
+python-multipart==0.0.16
+pytz==2024.2
+pywin32==308 ; platform_python_implementation != 'PyPy' and sys_platform == 'win32'
+pyyaml==6.0.2
+pyzmq==26.2.0
+requests==2.32.3
+rich==13.9.3 ; sys_platform != 'emscripten'
+ruff==0.7.1 ; sys_platform != 'emscripten'
+semantic-version==2.10.0
+setuptools==75.2.0 ; python_full_version >= '3.12'
+shellingham==1.5.4 ; sys_platform != 'emscripten'
+six==1.16.0
+sniffio==1.3.1
+stack-data==0.6.3
+starlette==0.41.2
+sympy==1.13.1
+tomlkit==0.12.0
+torch==2.5.0
+torchinfo==1.8.0
+torchvision==0.20.0
+tornado==6.4.1
+tqdm==4.66.5
+traitlets==5.14.3
+triton==3.1.0 ; python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'
+typer==0.12.5 ; sys_platform != 'emscripten'
+typing-extensions==4.12.2
+tzdata==2024.2
+urllib3==2.2.3
+uvicorn==0.32.0 ; sys_platform != 'emscripten'
+wcwidth==0.2.13
+websockets==12.0
+widgetsnbextension==4.0.13