File size: 4,669 Bytes
8c5a4e3 8969ba8 8c5a4e3 8969ba8 a365343 8c5a4e3 c96fce8 95f32ab 8c5a4e3 8385548 8c5a4e3 c96fce8 8c5a4e3 c96fce8 | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | import spaces # MUST come before torch / any CUDA-touching import
import random
import time
import torch
import gradio as gr
from inference import ShellDInference
MODEL_ID = "FlameF0X/ShellD"
TEXT_ENCODER = "sentence-transformers/all-MiniLM-L6-v2"
# Load at module scope, .to("cuda") eagerly. ShellDInference auto-detects cuda
# (the spaces hijack patches torch.cuda.is_available() in the main process).
# Pass the text encoder explicitly — the config.json's text_encoder_name is a
# local path ("./all-MiniLM-L6-v2") which doesn't exist in the Space.
pipe = ShellDInference(MODEL_ID, device="cpu", text_encoder_path=TEXT_ENCODER)
pipe.model.eval()
print("ShellD model loaded.")
def _estimate_duration(prompt, seed, randomize_seed, num_steps, cfg_scale, *args, **kwargs):
"""GPU-seconds reservation for ZeroGPU. Measured on ZeroGPU: ~0.01 s/step
for this 67M-param DiT (2 CFG forward passes per step on 64 patches).
Add ~5 s for cold-start weight streaming, floor at 10 s, ceiling at 60 s."""
return min(60, max(10, int(num_steps * 0.015 + 5)))
@spaces.GPU(duration=_estimate_duration)
def generate(
prompt: str,
seed: int = 42,
randomize_seed: bool = False,
num_steps: int = 250,
cfg_scale: float = 3.0,
progress=gr.Progress(track_tqdm=False),
):
"""Generate a 256×256 image from a text prompt using ShellD.
Args:
prompt: Text description of the image to generate.
seed: RNG seed for reproducibility (ignored if randomize_seed is True).
randomize_seed: If True, pick a random seed and write it back.
num_steps: Number of DDPM denoising steps (more = higher quality, slower).
cfg_scale: Classifier-free guidance scale (higher = more prompt adherence).
Returns:
(image, seed) — a 256×256 PIL Image and the seed used.
"""
if prompt is None or str(prompt).strip() == "":
prompt = "a serene lake surrounded by mountains"
prompt = str(prompt).strip()
if randomize_seed:
seed = random.randint(0, 2**31 - 1)
seed = int(seed)
t0 = time.perf_counter()
img = pipe.generate(
prompt=prompt,
num_steps=int(num_steps),
cfg_scale=float(cfg_scale),
seed=seed,
)
elapsed = time.perf_counter() - t0
print(f"ShellD generated '{prompt[:40]}' in {elapsed:.2f}s ({num_steps} steps, seed={seed})")
return img, seed
CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""
with gr.Blocks(css=CSS) as demo:
gr.Markdown(
"""
# 🐚 ShellD — Shell Diffusion Text-to-Image
A lightweight 67M-parameter Diffusion Transformer (DiT) that generates 256×256 images from text prompts.
[Model card](https://huggingface.co/FlameF0X/ShellD)
"""
)
with gr.Column(elem_id="col-container"):
with gr.Row():
prompt = gr.Textbox(
label="Prompt",
show_label=False,
placeholder="Describe the image you want to generate…",
container=False,
scale=4,
)
run = gr.Button("Generate", variant="primary", scale=1)
output = gr.Image(label="Generated image", height=320, show_label=True)
with gr.Accordion("Advanced settings", open=False):
num_steps = gr.Slider(
label="Steps", minimum=10, maximum=1000, step=10, value=250
)
cfg_scale = gr.Slider(
label="CFG scale", minimum=1.0, maximum=10.0, step=0.5, value=3.0
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
seed = gr.Number(label="Seed", value=0, precision=0)
gr.Examples(
examples=[
["a serene lake surrounded by mountains"],
["a futuristic city at night with neon lights"],
["a cute cat sitting on a windowsill"],
["a tropical beach with palm trees at sunset"],
],
inputs=[prompt],
outputs=[output, seed],
fn=generate,
cache_examples=True,
cache_mode="lazy",
)
run.click(
generate,
inputs=[prompt, seed, randomize_seed, num_steps, cfg_scale],
outputs=[output, seed],
api_name="generate",
)
prompt.submit(
generate,
inputs=[prompt, seed, randomize_seed, num_steps, cfg_scale],
outputs=[output, seed],
api_name="generate_submit",
)
if __name__ == "__main__":
demo.launch(mcp_server=True, theme=gr.themes.Citrus()) |