Spaces:
Running on Zero
Running on Zero
| """Graphics Studio on ZeroGPU. | |
| Image tab (text to image, image generation): FLUX.2 klein 9B turns a prompt into a still image. | |
| Music tab (text to music, text to audio): MusicGen medium turns a description into up to 28 seconds of music. | |
| Transcribe tab (speech to text, speech recognition, automatic speech recognition): Whisper large v3 turbo | |
| writes down speech from an uploaded file or the microphone, in its own language or translated to English. | |
| Narration tab (text to speech, machine speech, speech generation): Kokoro reads your text aloud. | |
| Describe tab: three sub-tabs for the describing jobs the Space does. | |
| Image to text (image captioning, image recognition): Florence-2 large writes a detailed caption for | |
| an uploaded picture. | |
| Video to text (video captioning, video to text, video-text-to-text): SmolVLM2 2.2B looks at frames | |
| from an uploaded clip and describes the scene and the motion. | |
| Music to text (music captioning, audio captioning, music classification, audio classification): | |
| CLAP picks a genre, mood, tempo and instruments for a piece of music and writes them as one short | |
| sentence. | |
| Video tab (text to video, image to video, video generation): one button makes a whole clip. A small | |
| text model (Qwen3 4B) first improves the scene and motion prompts if that is ticked, FLUX draws | |
| the scene, LTX-Video (13B) animates it, MusicGen medium adds a soundtrack if you describe one, | |
| and if narration is ticked Qwen3 4B writes the lines and their timing and Kokoro speaks them. | |
| Each model has its own ZeroGPU pack (see zero_packs.py), and each button copies only the packs it | |
| needs onto the GPU, one after another where a job uses several big models. That keeps every job | |
| inside the default half-card slice, which is not billed double. If zero_packs is switched off or | |
| cannot work, everything is copied on every request as ZeroGPU normally does, on the full card. | |
| """ | |
| import os | |
| def _pick_model_store(): | |
| """Off unless MODEL_STORE is set. Tried /data (the attached bucket) on 2026-10-04: it reads at | |
| about 8 MB/s, so loading 70 GB of models from it takes hours. The local disk downloads at 1 GB/s.""" | |
| path = os.environ.get("MODEL_STORE") | |
| if not path: | |
| return None | |
| try: | |
| os.makedirs(os.path.join(path, "huggingface"), exist_ok=True) | |
| return os.path.join(path, "huggingface") | |
| except OSError: | |
| return None | |
| _store = _pick_model_store() | |
| if _store: | |
| os.environ["HF_HOME"] = _store # Must be set before the Hugging Face libraries are imported | |
| import json | |
| import random | |
| import re | |
| import shutil | |
| import subprocess | |
| import sys | |
| import threading | |
| import time | |
| import gradio as gr | |
| import imageio | |
| import imageio_ffmpeg | |
| import numpy as np | |
| import scipy.io.wavfile | |
| import spaces | |
| import torch | |
| import zero_packs as packs | |
| from diffusers import FluxPipeline, LTXConditionPipeline | |
| from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition | |
| from diffusers.utils import export_to_video | |
| from PIL import Image | |
| from transformers import (AutoModelForCausalLM, AutoModelForImageTextToText, AutoProcessor, | |
| AutoTokenizer, MusicgenForConditionalGeneration, pipeline) | |
| try: | |
| from transformers import AutoModelForVision2Seq | |
| except ImportError: | |
| AutoModelForVision2Seq = None | |
| IMAGE_MODEL = "black-forest-labs/FLUX.2-klein-9B" # Gated, FLUX non-commercial licence | |
| IMAGE_FALLBACK = "black-forest-labs/FLUX.1-schnell" # Used if klein cannot be loaded (Apache 2.0) | |
| VIDEO_MODEL = "Lightricks/LTX-Video-0.9.8-13B-distilled" | |
| MUSIC_MODEL = "facebook/musicgen-medium" # Non-commercial licence (CC-BY-NC 4.0) | |
| TEXT_MODEL = "Qwen/Qwen3-4B" # Narration planner and prompt rewriter | |
| IMAGE_CAPTION_MODEL = "florence-community/Florence-2-large" # Picture to caption | |
| VIDEO_CAPTION_MODEL = "HuggingFaceTB/SmolVLM2-2.2B-Instruct" # Video frames to description | |
| LISTEN_MODEL = "laion/clap-htsat-unfused" # Recognises genre, mood, instruments and tempo | |
| SPEECH_MODEL = "openai/whisper-large-v3-turbo" # Speech to text | |
| VOICE = "af_heart" # Kokoro voice | |
| VOICE_RATE = 24000 | |
| # All 28 English voices of Kokoro. The first eight keep the names they had when there were only eight. | |
| VOICES = {"Heart (woman)": "af_heart", "Bella (woman)": "af_bella", "Nicole (woman, soft)": "af_nicole", | |
| "Sarah (woman)": "af_sarah", "Adam (man)": "am_adam", "Michael (man)": "am_michael", | |
| "Fenrir (man, deep)": "am_fenrir", "Puck (man, lively)": "am_puck", | |
| "Alloy (woman)": "af_alloy", "Aoede (woman)": "af_aoede", "Jessica (woman)": "af_jessica", | |
| "Kore (woman)": "af_kore", "Nova (woman)": "af_nova", "River (woman)": "af_river", "Sky (woman)": "af_sky", | |
| "Echo (man)": "am_echo", "Eric (man)": "am_eric", "Liam (man)": "am_liam", "Onyx (man)": "am_onyx", | |
| "Santa (man)": "am_santa", | |
| "Alice (British woman)": "bf_alice", "Emma (British woman)": "bf_emma", | |
| "Isabella (British woman)": "bf_isabella", "Lily (British woman)": "bf_lily", | |
| "Daniel (British man)": "bm_daniel", "Fable (British man)": "bm_fable", | |
| "George (British man)": "bm_george", "Lewis (British man)": "bm_lewis"} | |
| MAX_NARRATION_CHARS = 3000 | |
| WORDS_PER_SECOND = 2.5 | |
| IMAGE_WIDTH, IMAGE_HEIGHT = 1152, 768 # 3:2, multiples of 16 for FLUX | |
| VIDEO_WIDTH, VIDEO_HEIGHT = 1152, 768 # Same size as the image, so nothing is shrunk (multiples of 32) | |
| VIDEO_QUALITY = 9 # 0 to 10. Low values make the saved file blurry | |
| FPS = 10 | |
| TRIM_FRAMES = 8 # Frames cut from the start of every clip, where the noise is | |
| MAX_SEED = 2**31 - 1 | |
| NEGATIVE_PROMPT = ("worst quality, inconsistent motion, blurry, blurry motion, motion blur, " | |
| "ignoring the prompt, jittery, distorted, watermark, text, deformed") | |
| NEGATIVE_SLOW = ", slow motion, slow-motion video" | |
| STYLE_SUFFIX = "Sharp, crisp, highly detailed, in focus, no blur" | |
| SPEED_SUFFIX = ", moving at natural real-time speed" | |
| # ZeroGPU attaches the GPU only inside @spaces.GPU calls. Each model is sent "to cuda" here and then | |
| # sealed into its own pack, so a call copies only the packs it asks for with packs.load(...). | |
| # Pure per-call .to('cuda') from ordinary RAM was tried on 2026-10-07: the first FLUX move took 79 s. | |
| # The packs are read straight from disk at about 17 GB a second, like ZeroGPU's own single pack. | |
| # Once a model is sealed, never call .to("cuda") on it again here: its weights are on disk. | |
| print("[load] image model...", flush=True) | |
| try: | |
| from diffusers import Flux2KleinPipeline | |
| image_pipe = Flux2KleinPipeline.from_pretrained(IMAGE_MODEL, torch_dtype=torch.bfloat16) | |
| IMAGE_IS_KLEIN = True | |
| except Exception as e: # Licence not accepted on the account, or diffusers too old: the Space still starts | |
| print(f"[load] {IMAGE_MODEL} failed ({type(e).__name__}: {str(e)[:200]}); using {IMAGE_FALLBACK}", flush=True) | |
| image_pipe = FluxPipeline.from_pretrained(IMAGE_FALLBACK, torch_dtype=torch.bfloat16) | |
| IMAGE_IS_KLEIN = False | |
| for part in (image_pipe.transformer, image_pipe.vae, image_pipe.text_encoder, | |
| getattr(image_pipe, "text_encoder_2", None)): | |
| if part is not None: | |
| part.to("cuda") | |
| packs.seal("image", IMAGE_MODEL if IMAGE_IS_KLEIN else IMAGE_FALLBACK) | |
| print("[load] LTX...", flush=True) | |
| video_pipe = LTXConditionPipeline.from_pretrained(VIDEO_MODEL, torch_dtype=torch.bfloat16) | |
| video_pipe.vae.enable_tiling() # Decode the video in tiles to keep memory low | |
| video_pipe.text_encoder.to("cuda") | |
| packs.seal("t5") | |
| for part in (video_pipe.transformer, video_pipe.vae): | |
| part.to("cuda") | |
| packs.seal("ltx", VIDEO_MODEL) | |
| print("[load] MusicGen...", flush=True) | |
| music_processor = AutoProcessor.from_pretrained(MUSIC_MODEL) | |
| music_model = MusicgenForConditionalGeneration.from_pretrained( | |
| MUSIC_MODEL, torch_dtype=torch.float16).to("cuda") | |
| packs.seal("musicgen", MUSIC_MODEL) | |
| print("[load] Qwen3 4B...", flush=True) | |
| text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL) | |
| text_model = AutoModelForCausalLM.from_pretrained(TEXT_MODEL, torch_dtype=torch.bfloat16).to("cuda") | |
| packs.seal("qwen", TEXT_MODEL) | |
| print("[load] Florence-2...", flush=True) | |
| image_caption_processor = AutoProcessor.from_pretrained(IMAGE_CAPTION_MODEL, trust_remote_code=True) | |
| # florence-community is native; AutoModelForVision2Seq / AutoModelForImageTextToText register Florence2Config. | |
| _florence_err = None | |
| image_caption_model = None | |
| for _cls in (AutoModelForVision2Seq, AutoModelForImageTextToText, AutoModelForCausalLM): | |
| if _cls is None: | |
| continue | |
| try: | |
| image_caption_model = _cls.from_pretrained( | |
| IMAGE_CAPTION_MODEL, torch_dtype=torch.bfloat16, trust_remote_code=True).to("cuda") | |
| print(f"[load] Florence-2 via {_cls.__name__}", flush=True) | |
| break | |
| except Exception as e: | |
| _florence_err = e | |
| print(f"[load] Florence-2 via {_cls.__name__} failed: {type(e).__name__}: {e}", flush=True) | |
| if image_caption_model is None: | |
| raise RuntimeError(f"Could not load Florence-2: {_florence_err}") | |
| packs.seal("florence", IMAGE_CAPTION_MODEL) | |
| print("[load] SmolVLM2...", flush=True) | |
| video_caption_processor = AutoProcessor.from_pretrained(VIDEO_CAPTION_MODEL) | |
| video_caption_model = AutoModelForImageTextToText.from_pretrained( | |
| VIDEO_CAPTION_MODEL, torch_dtype=torch.bfloat16).to("cuda") | |
| packs.seal("smolvlm", VIDEO_CAPTION_MODEL) | |
| print("[load] CLAP...", flush=True) | |
| listen_pipe = pipeline("zero-shot-audio-classification", model=LISTEN_MODEL, device="cuda") | |
| packs.seal("clap", LISTEN_MODEL) | |
| print("[load] Whisper...", flush=True) | |
| # No chunk_length_s: recordings longer than 30 seconds are then read with Whisper's own long-form | |
| # method, one window after another. The pipeline's chunking was tried on 2026-10-11 and gave pieces | |
| # with wrong times, repeated text and missing passages. | |
| speech_pipe = pipeline("automatic-speech-recognition", model=SPEECH_MODEL, torch_dtype=torch.float16, | |
| device="cuda") | |
| packs.seal("whisper", SPEECH_MODEL) | |
| print("[load] done.", flush=True) | |
| # With a pack per model every job fits the default half-card slice ("large"), which is billed once. | |
| # Without them ZeroGPU copies all the models on every call, and that needs the whole card. | |
| # VIDEO_GPU_SIZE can be set to "xlarge" in the Space's variables if the video jobs run out of memory. | |
| GPU_SIZE = "large" if packs.enabled else "xlarge" | |
| VIDEO_GPU_SIZE = os.environ.get("VIDEO_GPU_SIZE", GPU_SIZE) if packs.enabled else "xlarge" | |
| def gpu_seconds(seconds): | |
| """The GPU time a job asks for. The numbers passed in were measured on the half card on | |
| 2026-10-08, where ZeroGPU counts 1.5 times what is asked (about 20% headroom is included). | |
| A job that runs past its time is cut off. Without per-model packs every job also waits about | |
| 5 s for all the models to be copied, and the full card counts double.""" | |
| return int(seconds) if packs.enabled else int(seconds + 6) | |
| # ---------------------------------------------------------------- Timing helper | |
| class timed: | |
| """Prints how long a block of code took.""" | |
| def __init__(self, label): | |
| self.label = label | |
| def __enter__(self): | |
| self.t0 = time.time() | |
| return self | |
| def __exit__(self, *_): | |
| print(f"[{self.label}] {time.time() - self.t0:.2f}s", flush=True) | |
| # ---------------------------------------------------------------- Kokoro (runs on the CPU in its own process) | |
| def start_voice_worker(): | |
| worker = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tts_worker.py") | |
| return subprocess.Popen([sys.executable, worker], stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True) | |
| voice_lock = threading.Lock() | |
| voice_worker = start_voice_worker() | |
| def kokoro(text, voice, speed=1.0): | |
| """The text spoken by Kokoro, as a float32 wave at VOICE_RATE.""" | |
| global voice_worker | |
| with voice_lock: | |
| if voice_worker.poll() is not None: | |
| voice_worker = start_voice_worker() | |
| out = f"/tmp/kokoro_{random.randint(0, MAX_SEED)}.npy" | |
| voice_worker.stdin.write(json.dumps({"text": text, "voice": voice, "speed": speed, "out": out}) + "\n") | |
| voice_worker.stdin.flush() | |
| answer = voice_worker.stdout.readline() | |
| if not answer: | |
| raise gr.Error("The narration model stopped. Please try again.") | |
| answer = json.loads(answer) | |
| if not answer["ok"]: | |
| raise gr.Error("Narration failed: " + answer["error"]) | |
| wave = np.load(out) | |
| os.remove(out) | |
| return wave | |
| def pick_seed(seed, randomize): | |
| return random.randint(0, MAX_SEED) if randomize else int(seed) | |
| # ---------------------------------------------------------------- The core model steps (run on the GPU) | |
| def draw_image(prompt, seed): | |
| # Both are 4-step distilled models; they only differ in how guidance is switched off | |
| options = {"guidance_scale": 1.0} if IMAGE_IS_KLEIN else {"guidance_scale": 0.0, "max_sequence_length": 256} | |
| return image_pipe( | |
| prompt=prompt, | |
| width=IMAGE_WIDTH, | |
| height=IMAGE_HEIGHT, | |
| num_inference_steps=4, | |
| generator=torch.Generator("cpu").manual_seed(seed), | |
| **options, | |
| ).images[0] | |
| def animate(image, scene, motion, seconds, seed): | |
| """Returns the path of a silent clip and its length in seconds.""" | |
| num_frames = clip_frames(seconds) | |
| motion = motion.strip() or "gentle cinematic camera motion, natural movement" | |
| wants_slow = re.search(r"slow[\s-]?mo", motion.lower()) is not None | |
| style = STYLE_SUFFIX + ("" if wants_slow else SPEED_SUFFIX) + "." | |
| prompt = f"{motion}. {scene.strip()}. {style}" if scene.strip() else f"{motion}. {style}" | |
| start = LTXVideoCondition(image=image.convert("RGB").resize((VIDEO_WIDTH, VIDEO_HEIGHT)), frame_index=0) | |
| frames = video_pipe( | |
| conditions=[start], | |
| prompt=prompt, | |
| negative_prompt=NEGATIVE_PROMPT + ("" if wants_slow else NEGATIVE_SLOW), | |
| image_cond_noise_scale=0.0, | |
| decode_timestep=0.05, | |
| decode_noise_scale=0.025, | |
| width=VIDEO_WIDTH, | |
| height=VIDEO_HEIGHT, | |
| num_frames=num_frames, | |
| frame_rate=FPS, | |
| timesteps=[1000, 993, 987, 981, 975, 909, 725, 0.03], | |
| guidance_scale=1.0, | |
| guidance_rescale=0.7, | |
| generator=torch.Generator("cpu").manual_seed(seed), | |
| ).frames[0] | |
| frames = frames[TRIM_FRAMES:] | |
| path = f"/tmp/video_{seed}_{random.randint(0, 99999)}.mp4" | |
| export_to_video(frames, path, fps=FPS, quality=VIDEO_QUALITY) | |
| return path, len(frames) / FPS | |
| def compose_music(music_prompt, seconds): | |
| """Returns the path of a WAV file as long as the clip.""" | |
| seconds = min(max(seconds, 2.0), 28.0) | |
| inputs = music_processor(text=[music_prompt], padding=True, return_tensors="pt").to("cuda") | |
| audio = music_model.generate(**inputs, do_sample=True, guidance_scale=3.0, | |
| max_new_tokens=int(seconds * 50) + 8) | |
| rate = music_model.config.audio_encoder.sampling_rate | |
| wave = audio[0, 0].float().cpu().numpy()[:int(seconds * rate)] | |
| fade = int(0.25 * rate) | |
| wave[:fade] *= np.linspace(0.0, 1.0, fade) | |
| wave[-fade:] *= np.linspace(1.0, 0.0, fade) | |
| wav_path = f"/tmp/music_{random.randint(0, MAX_SEED)}.wav" | |
| scipy.io.wavfile.write(wav_path, rate, (np.clip(wave, -1, 1) * 32767).astype(np.int16)) | |
| return wav_path | |
| # ---------------------------------------------------------------- Describing (vision / audio) and rewriting (text) | |
| GENRES = ["classical orchestral", "electronic dance", "rock", "jazz", "hip hop", "ambient", "pop", "folk acoustic", | |
| "8-bit chiptune", "cinematic soundtrack", "lo-fi", "heavy metal", "Indian classical", "devotional chant"] | |
| MOODS = ["happy and upbeat", "calm and peaceful", "sad and emotional", "tense and dramatic", "epic and heroic", | |
| "dark and mysterious", "playful and fun", "dreamy and floating"] | |
| INSTRUMENTS = ["piano", "acoustic guitar", "electric guitar", "violins and strings", "drums", "synthesizer", "flute", | |
| "bass", "brass and trumpets", "singing voice", "bells", "sitar", "tabla"] | |
| TEMPOS = ["slow tempo", "medium tempo", "fast tempo"] | |
| # Chat models like to open with "Certainly! Here's a revised prompt:". That opening is cut off, | |
| # because the image, video and music models would otherwise be given it as part of the prompt. | |
| OPENING_WORD = re.compile(r"^(certainly|sure|of course|absolutely|okay|ok|great)\b[!,.]?\s*", re.IGNORECASE) | |
| OPENING_LINE = re.compile(r"^(here(['’]s| is| are)|below is|this is)\b[^\n.!?]*:\s*", re.IGNORECASE) | |
| def tidy(text, one_line=False): | |
| """The model's answer without its chatty first sentence and without Markdown marks.""" | |
| text = text.strip() | |
| text = OPENING_WORD.sub("", text, count=1) | |
| text = OPENING_LINE.sub("", text, count=1) | |
| text = re.sub(r"^#+\s*", "", text, flags=re.MULTILINE) | |
| text = re.sub(r"[*`]+", "", text).strip().strip('"').strip() | |
| return " ".join(text.split()) if one_line else text | |
| def shrink(image, side=768): | |
| image = image.convert("RGB") | |
| image.thumbnail((side, side)) | |
| return image | |
| def caption_image(image): | |
| """Florence-2 writes a detailed caption. Model must already be on the GPU.""" | |
| prompt = "<MORE_DETAILED_CAPTION>" | |
| inputs = image_caption_processor(text=prompt, images=image, return_tensors="pt").to("cuda", torch.bfloat16) | |
| out = image_caption_model.generate( | |
| input_ids=inputs["input_ids"], | |
| pixel_values=inputs["pixel_values"], | |
| max_new_tokens=512, | |
| num_beams=3, | |
| do_sample=False, | |
| ) | |
| decoded = image_caption_processor.batch_decode(out, skip_special_tokens=False)[0] | |
| parsed = image_caption_processor.post_process_generation( | |
| decoded, task=prompt, image_size=(image.width, image.height)) | |
| return parsed.get(prompt, "").strip() | |
| def caption_video_frames(frames): | |
| """SmolVLM2 describes a short video from a few frames. Model must already be on the GPU.""" | |
| content = [{"type": "image", "image": f} for f in frames] | |
| content.append({"type": "text", "text": "These pictures are frames from one video, in order. " | |
| "Describe the scene, then describe what moves and how."}) | |
| inputs = video_caption_processor.apply_chat_template( | |
| [{"role": "user", "content": content}], | |
| add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt").to("cuda") | |
| if "pixel_values" in inputs: | |
| inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16) | |
| out = video_caption_model.generate(**inputs, max_new_tokens=300, do_sample=False) | |
| text = video_caption_processor.batch_decode( | |
| out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0] | |
| return tidy(text) | |
| def caption_music(wave): | |
| """CLAP picks a genre, mood, tempo and two instruments, written as one short sentence. | |
| Model must already be on the GPU.""" | |
| def best(labels, count=1): | |
| return [hit["label"] for hit in listen_pipe(wave, candidate_labels=labels)[:count]] | |
| return (f"{best(GENRES)[0]} music, {best(MOODS)[0]}, {best(TEMPOS)[0]}, " | |
| f"featuring {' and '.join(best(INSTRUMENTS, 2))}") | |
| def rewrite(system, user, max_new_tokens=160, one_line=True): | |
| """Qwen3 4B edits a description into a better prompt. Model must already be on the GPU.""" | |
| text = text_tokenizer.apply_chat_template( | |
| [{"role": "system", "content": system}, {"role": "user", "content": user}], | |
| tokenize=False, add_generation_prompt=True, enable_thinking=False) | |
| inputs = text_tokenizer([text], return_tensors="pt").to("cuda") | |
| out = text_model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=True, | |
| temperature=0.6, top_p=0.9) | |
| return tidy(text_tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True), one_line) | |
| def plan_narration(scene, motion, seconds): | |
| """Qwen3 4B writes 1 to 3 lines of narration with start times. Model must already be on the GPU.""" | |
| budget = max(4, int((seconds - 1) * WORDS_PER_SECOND)) | |
| system = ("You are the narrator of a short film clip. You are given what the clip shows. Write what the " | |
| "narrator SAYS over it: a thought, a feeling or a tiny story about the scene, like a nature " | |
| "documentary. Never repeat or list the description itself, and never mention cameras, prompts or " | |
| "video. Reply with 1 to 3 lines and nothing else. Each line is: start second | sentence.\n" | |
| f"The clip lasts {seconds:.0f} seconds. Use at most {budget} words in total, so every line is short. " | |
| "The first line starts at 0.5 or later, and later lines start after the earlier ones have been spoken.\n\n" | |
| "Example for a clip of a lighthouse in a storm, 6 seconds:\n" | |
| "0.5 | Some lights refuse to go out.\n" | |
| "3.5 | Even when the sea says otherwise.") | |
| user = f"The clip shows: {scene or 'a quiet scene'}\nWhat moves: {motion or 'gentle movement'}" | |
| text = text_tokenizer.apply_chat_template( | |
| [{"role": "system", "content": system}, {"role": "user", "content": user}], | |
| tokenize=False, add_generation_prompt=True, enable_thinking=False) | |
| inputs = text_tokenizer([text], return_tensors="pt").to("cuda") | |
| out = text_model.generate(**inputs, max_new_tokens=120, do_sample=True, temperature=0.8, top_p=0.9) | |
| return text_tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) | |
| # ---------------------------------------------------------------- Narration and mixing (CPU only) | |
| def parse_plan(reply, seconds): | |
| """Turn the planner's reply into [(start, text)]. If it ignored the format, its own words are | |
| still used as one line; the scene description is never read out.""" | |
| plan = [] | |
| for row in reply.splitlines(): | |
| match = re.match(r"\s*[-*]?\s*(\d+(?:\.\d+)?)\s*(?:s|sec|seconds)?\s*[|:–-]\s*(.+)", row) | |
| if match and match.group(2).strip(): | |
| plan.append((float(match.group(1)), match.group(2).strip().strip('"'))) | |
| if not plan: | |
| words = re.sub(r"[\[\]{}\"|]", " ", reply).split()[:int(seconds * WORDS_PER_SECOND)] | |
| plan = [(0.5, " ".join(words) or "A quiet moment, captured in motion.")] | |
| return sorted(plan)[:3] | |
| def speak(text): | |
| return kokoro(text, VOICE) | |
| def narrate(reply, seconds): | |
| """Speak the plan into one WAV as long as the clip. Returns (wav path or None, the plan as text).""" | |
| track = np.zeros(int(seconds * VOICE_RATE), dtype=np.float32) | |
| cursor, script = 0.0, [] | |
| for start, line in parse_plan(reply, seconds): | |
| voice = speak(line) | |
| start = max(start, cursor) | |
| begin = int(start * VOICE_RATE) | |
| room = len(track) - begin | |
| if room < VOICE_RATE // 2: | |
| continue | |
| voice = voice[:room] | |
| track[begin:begin + len(voice)] += voice | |
| cursor = start + len(voice) / VOICE_RATE + 0.3 | |
| script.append(f"{start:4.1f} s {line}") | |
| if not script: | |
| return None, "The clip is too short for narration." | |
| path = f"/tmp/voice_{random.randint(0, MAX_SEED)}.wav" | |
| scipy.io.wavfile.write(path, VOICE_RATE, (np.clip(track, -1, 1) * 32767).astype(np.int16)) | |
| return path, "\n".join(script) | |
| def mix(video, music_wav, voice_wav): | |
| """Put whichever sound tracks exist onto the silent clip. With both, music plays under the voice.""" | |
| if not music_wav and not voice_wav: | |
| return video | |
| ffmpeg = [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-i", video] | |
| if music_wav and voice_wav: | |
| ffmpeg += ["-i", voice_wav, "-i", music_wav, "-filter_complex", | |
| "[1:a]aresample=44100,volume=1.6[v];[2:a]aresample=44100,volume=0.35[m];" | |
| "[v][m]amix=inputs=2:duration=longest:normalize=0[a]", "-map", "0:v:0", "-map", "[a]"] | |
| else: | |
| ffmpeg += ["-i", music_wav or voice_wav, "-map", "0:v:0", "-map", "1:a:0"] | |
| out_path = f"/tmp/mixed_{random.randint(0, MAX_SEED)}.mp4" | |
| subprocess.run(ffmpeg + ["-c:v", "copy", "-c:a", "aac", "-shortest", out_path], check=True, capture_output=True) | |
| return out_path | |
| def load_audio(path, limit=30): | |
| """Any audio file as mono 48 kHz samples (the first 30 seconds), using ffmpeg.""" | |
| wav = f"/tmp/listen_{random.randint(0, MAX_SEED)}.wav" | |
| subprocess.run([imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-i", path, "-t", str(limit), "-ac", "1", "-ar", "48000", wav], | |
| check=True, capture_output=True) | |
| rate, data = scipy.io.wavfile.read(wav) | |
| return data.astype(np.float32) / 32768.0 | |
| def video_frames(path, count=6): | |
| """A few frames spread across the clip, and its length in seconds.""" | |
| reader = imageio.get_reader(path) | |
| meta = reader.get_meta_data() | |
| length = float(meta.get("duration") or 5.0) | |
| total = max(1, int(length * float(meta.get("fps") or 10)) - 1) | |
| frames = [] | |
| for k in range(count): | |
| try: | |
| frames.append(shrink(Image.fromarray(reader.get_data(total * k // max(1, count - 1))), 512)) | |
| except Exception: | |
| break | |
| reader.close() | |
| return frames, length | |
| # ---------------------------------------------------------------- The buttons | |
| def make_image(prompt, seed, randomize_seed): | |
| if not prompt.strip(): | |
| raise gr.Error("Describe the image first.") | |
| seed = pick_seed(seed, randomize_seed) | |
| with timed("image"): | |
| packs.load("image") | |
| return draw_image(prompt, seed), seed | |
| def music_duration(prompt, seconds): | |
| return gpu_seconds(3 + seconds * 1.0) | |
| def make_music(prompt, seconds): | |
| if not prompt.strip(): | |
| raise gr.Error("Describe the music first.") | |
| with timed("music"): | |
| packs.load("musicgen") | |
| return compose_music(prompt, seconds) | |
| def audio_seconds(path): | |
| """Length of an audio file in seconds, read from ffmpeg's report. 30 if it cannot be read.""" | |
| try: | |
| report = subprocess.run([imageio_ffmpeg.get_ffmpeg_exe(), "-hide_banner", "-i", path], | |
| capture_output=True, text=True, timeout=30).stderr | |
| hours, minutes, seconds = re.search(r"Duration: (\d+):(\d+):(\d+\.?\d*)", report).groups() | |
| return int(hours) * 3600 + int(minutes) * 60 + float(seconds) | |
| except Exception: | |
| return 30.0 | |
| def transcribe_duration(audio, task): | |
| # A short clip needs about 3 s in all. Long recordings are read in 30-second pieces; this allows | |
| # for the model working at 8 times real speed, which is slower than it has measured so far. | |
| return gpu_seconds(5 + audio_seconds(audio) / 12) if audio else gpu_seconds(5) | |
| def transcribe(audio, task): | |
| if not audio: | |
| raise gr.Error("Upload a recording or record with the microphone first.") | |
| with timed("transcribe"): | |
| packs.load("whisper") | |
| options = {"task": "translate" if task.startswith("Translate") else "transcribe"} | |
| if audio_seconds(audio) > 30: # Long-form only: retry a window that came out badly, skip silence | |
| options.update(temperature=(0.0, 0.2, 0.4, 0.6, 0.8, 1.0), no_speech_threshold=0.6, | |
| logprob_threshold=-1.0, compression_ratio_threshold=1.35) | |
| result = speech_pipe(audio, return_timestamps=True, generate_kwargs=options) | |
| # The same text again, one piece per line with its start and end time in seconds | |
| timed_lines = [] | |
| for piece in result.get("chunks") or []: | |
| start, end = piece.get("timestamp") or (None, None) | |
| if start is not None and piece.get("text", "").strip(): | |
| timed_lines.append(f"[{start:.2f} - {end if end is not None else start:.2f}] {piece['text'].strip()}") | |
| return result["text"].strip(), "\n".join(timed_lines) | |
| def read_aloud(text, voice_name, speed): | |
| """Kokoro runs on the CPU, so this uses no GPU quota.""" | |
| text = text.strip() | |
| if not text: | |
| raise gr.Error("Type the text to read first.") | |
| if len(text) > MAX_NARRATION_CHARS: | |
| raise gr.Error(f"That is {len(text)} characters. Please keep it under {MAX_NARRATION_CHARS}.") | |
| wave = kokoro(text, VOICES[voice_name], float(speed)) | |
| if not len(wave): | |
| raise gr.Error("Nothing could be read from that text.") | |
| path = f"/tmp/narration_{random.randint(0, MAX_SEED)}.wav" | |
| scipy.io.wavfile.write(path, VOICE_RATE, (np.clip(wave, -1, 1) * 32767).astype(np.int16)) | |
| return path | |
| # ---------------------------------------------------------------- Describe tab | |
| def describe_image_tab(image): | |
| if image is None: | |
| raise gr.Error("Upload an image first.") | |
| small = shrink(image) | |
| with timed("image-to-text"): | |
| packs.load("florence") | |
| return caption_image(small) | |
| def describe_video_tab(video): | |
| if not video: | |
| raise gr.Error("Upload a video first.") | |
| frames, _ = video_frames(video) | |
| if not frames: | |
| raise gr.Error("That video could not be read.") | |
| with timed("video-to-text"): | |
| packs.load("smolvlm") | |
| return caption_video_frames(frames) | |
| def describe_music_tab(audio): | |
| if not audio: | |
| raise gr.Error("Upload some music first.") | |
| try: | |
| wave = load_audio(audio) | |
| except Exception: | |
| raise gr.Error("That audio file could not be read.") | |
| with timed("music-to-text"): | |
| packs.load("clap") | |
| return caption_music(wave) | |
| # ---------------------------------------------------------------- Video tab | |
| def improve_prompts(scene, motion): | |
| """Qwen3 4B adds visual detail to the scene and motion prompts without changing what they ask for.""" | |
| reply = rewrite( | |
| "You improve prompts for a video generator. You are given a scene and a motion. Keep every subject, " | |
| "action and direction exactly as given (left stays left, right stays right), and never add new subjects " | |
| "or change what moves. Add what a camera would see: lighting, colours, textures, atmosphere, depth. " | |
| "Reply with exactly two lines and nothing else:\nSCENE: the picture, under 60 words\n" | |
| "MOTION: what moves and how, under 30 words", | |
| f"Scene: {scene.strip() or 'not given, keep it short and neutral'}\n" | |
| f"Motion: {motion.strip() or 'gentle natural movement'}", 200, one_line=False) | |
| # The model answers in several shapes: "SCENE: ... MOTION: ..." on one line or two, or two bare lines | |
| label = re.search(r"\bmotion\s*:\s*", reply, re.IGNORECASE) | |
| lines = [line.strip() for line in reply.splitlines() if line.strip()] | |
| if label: | |
| new_scene, new_motion = reply[:label.start()], reply[label.end():] | |
| elif len(lines) == 2: | |
| new_scene, new_motion = lines | |
| else: | |
| new_scene = new_motion = "" | |
| new_scene = " ".join(re.sub(r"^\s*scene\s*:\s*", "", new_scene, flags=re.IGNORECASE).split()) | |
| new_motion = " ".join(new_motion.split()) | |
| if not new_scene or not new_motion: # Anything else: keep the visitor's own words | |
| print(f"[improve] unusable answer: {reply[:200]!r}", flush=True) | |
| return scene, motion | |
| return (new_scene if scene.strip() else scene), new_motion | |
| def clip_frames(seconds): | |
| """Frames LTX-Video is asked for: a multiple of 8 plus 1, with room for the trimmed start.""" | |
| return -(-(int(seconds * FPS) + TRIM_FRAMES - 1) // 8) * 8 + 1 | |
| def film(image, scene, motion, music_prompt, want_narration, improve, seconds, seed): | |
| """A clip made one model at a time, so FLUX and LTX-Video are never on the GPU together: | |
| the text model (prompts and narration plan), FLUX draws, LTX-Video animates, then MusicGen.""" | |
| plan = None | |
| improve = improve and bool(scene.strip() or motion.strip()) | |
| if improve or want_narration: | |
| packs.load("qwen") | |
| with timed(" text"): | |
| if improve: | |
| scene, motion = improve_prompts(scene, motion) | |
| if want_narration: | |
| plan = plan_narration(scene, motion, (clip_frames(seconds) - TRIM_FRAMES) / FPS) | |
| if image is None: | |
| packs.load("image") | |
| with timed(" draw"): | |
| image = draw_image(scene, seed) | |
| packs.load("t5", "ltx") | |
| with timed(" animate"): | |
| video, length = animate(image, scene, motion, seconds, seed) | |
| music_wav = None | |
| if music_prompt.strip(): | |
| packs.load("musicgen") | |
| with timed(" music"): | |
| music_wav = compose_music(music_prompt, length) | |
| return video, length, music_wav, plan, scene, motion | |
| def studio_duration(image, scene, motion, music_prompt, want_narration, improve, seconds, seed): | |
| total = 7 + seconds * 3.0 # Animation | |
| if image is None: | |
| total += 5 # FLUX draws the scene first | |
| if music_prompt.strip(): | |
| total += 2 + seconds * 1.0 | |
| if want_narration: | |
| total += 1 | |
| if improve: | |
| total += 2 # The text model rewrites the prompts | |
| return gpu_seconds(total) | |
| def run_models(image, scene, motion, music_prompt, want_narration, improve, seconds, seed): | |
| with timed("video"): | |
| return film(image, scene, motion, music_prompt, want_narration, improve, seconds, seed) | |
| def make_video(image, scene, motion, music_prompt, want_narration, improve, seconds, seed, randomize_seed): | |
| if image is None and not scene.strip(): | |
| raise gr.Error("Describe the scene, or add a start image.") | |
| seed = pick_seed(seed, randomize_seed) | |
| video, length, music_wav, plan, used_scene, used_motion = run_models( | |
| image, scene, motion, music_prompt, want_narration, improve, seconds, seed) | |
| voice_wav, script = narrate(plan, length) if plan is not None else (None, "") | |
| used = f"SCENE: {used_scene}\nMOTION: {used_motion}" if improve else "" | |
| return mix(video, music_wav, voice_wav), script, used, seed | |
| # ---------------------------------------------------------------- Bucket: running totals and showcase | |
| BUCKET = os.environ.get("STUDIO_BUCKET", "/data") | |
| COUNTS_FILE = os.path.join(BUCKET, "counts.json") | |
| SHOWCASE_DIR = os.path.join(BUCKET, "showcase") | |
| SHOWCASE_COPY = "/tmp/showcase" | |
| SHOWCASE_TYPES = (".png", ".jpg", ".jpeg", ".webp", ".mp4", ".webm") | |
| COUNT_NAMES = {"images": "images", "videos": "videos", "songs": "songs", "narrations": "narrations"} | |
| counts_lock = threading.Lock() | |
| counts = dict.fromkeys(COUNT_NAMES, 0) | |
| try: | |
| with open(COUNTS_FILE, encoding="utf-8") as f: | |
| counts.update({k: int(v) for k, v in json.load(f).items() if k in counts}) | |
| except (OSError, ValueError, AttributeError, TypeError): | |
| pass | |
| def counts_text(): | |
| total = sum(counts.values()) | |
| if not total: | |
| return "" | |
| parts = " · ".join(f"{counts[k]:,} {COUNT_NAMES[k][:-1] if counts[k] == 1 else COUNT_NAMES[k]}" | |
| for k in COUNT_NAMES if counts[k]) | |
| return f"**{total:,} creation{'' if total == 1 else 's'} so far:** {parts}" | |
| def bump(kind): | |
| with counts_lock: | |
| counts[kind] += 1 | |
| try: | |
| with open(COUNTS_FILE, "w", encoding="utf-8") as f: | |
| json.dump(counts, f) | |
| except OSError: | |
| pass | |
| return counts_text() | |
| def showcase(): | |
| try: | |
| names = sorted(n for n in os.listdir(SHOWCASE_DIR) if n.lower().endswith(SHOWCASE_TYPES)) | |
| except OSError: | |
| return [] | |
| os.makedirs(SHOWCASE_COPY, exist_ok=True) | |
| shown = [] | |
| for name in names: | |
| local = os.path.join(SHOWCASE_COPY, name) | |
| try: | |
| source = os.path.join(SHOWCASE_DIR, name) | |
| if not os.path.exists(local) or os.path.getsize(local) != os.path.getsize(source): # New, or replaced | |
| shutil.copy(source, local + ".part") | |
| os.replace(local + ".part", local) | |
| shown.append((local, os.path.splitext(name)[0].replace("_", " "))) | |
| except OSError: | |
| continue | |
| return shown | |
| def on_open(): | |
| return counts_text(), showcase() | |
| # ---------------------------------------------------------------- UI | |
| with gr.Blocks(title="Graphics Studio") as demo: | |
| gr.Markdown("# Graphics Studio\n**FLUX.2 klein** images, **LTX-Video** animation, **MusicGen** music, " | |
| "**Qwen3** + **Kokoro** narration, **Whisper** transcription, " | |
| "**Florence-2** image captions, **SmolVLM2** video descriptions, **CLAP** music classification.") | |
| gr.Markdown("Already have an image, a video or some music and want to change it? Try my other Space, " | |
| "[Graphics Enhancer](https://huggingface.co/spaces/vrishv/graphics_enhancer).") | |
| counter = gr.Markdown(counts_text()) | |
| with gr.Tab("Image"): | |
| image_prompt = gr.Textbox(label="Describe the image", lines=3, | |
| placeholder="A breathtaking aurora over a frozen lake at night, ultra detailed") | |
| with gr.Row(): | |
| image_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed") | |
| image_random = gr.Checkbox(value=True, label="Random seed") | |
| image_button = gr.Button("Generate image", variant="primary") | |
| image_out = gr.Image(label="Image", type="pil") | |
| gr.Examples( | |
| examples=[ | |
| ["A breathtaking aurora over a frozen lake at night, snowy mountains, ultra detailed"], | |
| ["A lone lighthouse on a cliff during a thunderstorm, giant waves, cinematic lighting"], | |
| ["A floating island with waterfalls at sunrise, golden clouds, fantasy landscape"], | |
| ], | |
| inputs=[image_prompt], | |
| ) | |
| with gr.Tab("Music"): | |
| song_prompt = gr.Textbox(label="Describe the music", lines=3, | |
| placeholder="Upbeat 8-bit arcade tune with a catchy melody, fast tempo") | |
| song_seconds = gr.Slider(5, 28, value=15, step=1, label="Length (seconds)") | |
| song_button = gr.Button("Generate music", variant="primary") | |
| song_out = gr.Audio(label="Music", type="filepath", autoplay=False) | |
| gr.Examples( | |
| examples=[ | |
| ["Upbeat 8-bit arcade tune with a catchy melody, fast tempo"], | |
| ["Calm ambient synth pads, slow and dreamy, night sky mood"], | |
| ["Epic orchestral trailer music with drums and brass, building tension"], | |
| ], | |
| inputs=[song_prompt], | |
| ) | |
| with gr.Tab("Transcribe"): | |
| speech_in = gr.Audio(label="Upload a recording, or speak now", sources=["upload", "microphone"], | |
| type="filepath") | |
| speech_task = gr.Radio(["Transcribe (keep the spoken language)", "Translate to English"], | |
| value="Transcribe (keep the spoken language)", label="What to write") | |
| speech_button = gr.Button("Transcribe", variant="primary") | |
| speech_out = gr.Textbox(label="Text", lines=8) | |
| speech_times = gr.Textbox(label="With times (start and end of each piece, in seconds)", lines=8, | |
| interactive=False) | |
| with gr.Tab("Narration"): | |
| read_text = gr.Textbox(label="Text to read aloud", lines=6, | |
| placeholder="Far above the sleeping town, the stars kept their quiet watch.") | |
| with gr.Row(): | |
| read_voice = gr.Dropdown(list(VOICES), value="Heart (woman)", label="Voice") | |
| read_speed = gr.Slider(0.7, 1.3, value=1.0, step=0.05, label="Speed") | |
| read_button = gr.Button("Read aloud", variant="primary") | |
| read_out = gr.Audio(label="Narration", type="filepath", autoplay=False) | |
| with gr.Tab("Describe"): | |
| gr.Markdown("Three ways to describe what you have: an image (image captioning), a short video " | |
| "(video captioning), or a piece of music (music captioning, music classification). " | |
| "The video and music jobs look at the first 30 seconds.") | |
| with gr.Tabs(): | |
| with gr.Tab("Image to text"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| desc_image_in = gr.Image(label="Your image", type="pil") | |
| desc_image_button = gr.Button("Describe image", variant="primary") | |
| with gr.Column(): | |
| desc_image_out = gr.Textbox(label="Description", lines=10, interactive=False) | |
| with gr.Tab("Video to text"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| desc_video_in = gr.Video(label="Your video", autoplay=False) | |
| desc_video_button = gr.Button("Describe video", variant="primary") | |
| with gr.Column(): | |
| desc_video_out = gr.Textbox(label="Description", lines=10, interactive=False) | |
| with gr.Tab("Music to text"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| desc_music_in = gr.Audio(label="Your music", sources=["upload"], type="filepath") | |
| desc_music_button = gr.Button("Describe music", variant="primary") | |
| with gr.Column(): | |
| desc_music_out = gr.Textbox(label="Description", lines=4, interactive=False) | |
| with gr.Tab("Video"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| scene_prompt = gr.Textbox(label="Scene", lines=3, | |
| placeholder="A lone lighthouse on a cliff during a thunderstorm, giant waves") | |
| motion_prompt = gr.Textbox(label="Motion", lines=2, | |
| placeholder="Waves crash against the rocks, the light beam sweeps to the right") | |
| music_prompt = gr.Textbox(label="Music (optional, leave empty for none)", lines=2, | |
| placeholder="Tense orchestral strings with distant thunder") | |
| improve_prompt = gr.Checkbox(value=True, label="Improve my prompts first (Qwen3 4B adds visual " | |
| "detail; untick to use your exact words)") | |
| want_narration = gr.Checkbox(value=False, label="Add narration") | |
| seconds = gr.Slider(2, 8, value=5, step=1, label="Length (seconds)") | |
| with gr.Accordion("More options", open=False): | |
| start_image = gr.Image(label="Start image (optional; otherwise FLUX draws the scene)", type="pil") | |
| with gr.Row(): | |
| video_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed") | |
| video_random = gr.Checkbox(value=True, label="Random seed") | |
| video_button = gr.Button("Make video", variant="primary") | |
| with gr.Column(): | |
| video_out = gr.Video(label="Video", autoplay=False, loop=True) | |
| narration_script = gr.Textbox(label="Narration (start time and line)", lines=3, interactive=False) | |
| used_prompts = gr.Textbox(label="Prompts used (after improving)", lines=4, interactive=False) | |
| with gr.Tab("Showcase"): | |
| gr.Markdown("A few favourites made with this Space.") | |
| showcase_out = gr.Gallery(label="Showcase", columns=3) | |
| image_button.click(make_image, [image_prompt, image_seed, image_random], [image_out, image_seed]).success( | |
| lambda: bump("images"), None, counter) | |
| song_button.click(make_music, [song_prompt, song_seconds], [song_out]).success( | |
| lambda: bump("songs"), None, counter) | |
| speech_button.click(transcribe, [speech_in, speech_task], [speech_out, speech_times]) | |
| read_button.click(read_aloud, [read_text, read_voice, read_speed], [read_out]).success( | |
| lambda: bump("narrations"), None, counter) | |
| desc_image_button.click(describe_image_tab, [desc_image_in], [desc_image_out]) | |
| desc_video_button.click(describe_video_tab, [desc_video_in], [desc_video_out]) | |
| desc_music_button.click(describe_music_tab, [desc_music_in], [desc_music_out]) | |
| video_button.click(make_video, | |
| [start_image, scene_prompt, motion_prompt, music_prompt, want_narration, improve_prompt, | |
| seconds, video_seed, video_random], | |
| [video_out, narration_script, used_prompts, video_seed]).success( | |
| lambda: bump("videos"), None, counter) | |
| demo.load(on_open, None, [counter, showcase_out]) | |
| demo.queue().launch(allowed_paths=[SHOWCASE_COPY]) | |