Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,429 @@
+"""
+DoodleDreams — ZeroGPU orchestrator
+Draw + voice → illustrated bedtime storybook narrated in your cloned voice.
+"""
+import os, sys, json, time, tempfile, logging, threading
+import torch
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+try:
+    import spaces
+except ModuleNotFoundError:
+    class _Shim:
+        @staticmethod
+        def GPU(*a, **k):
+            return a[0] if a and callable(a[0]) else (lambda fn: fn)
+    spaces = _Shim()
+
+import gradio as gr
+from config import (
+    FLUX_MODEL, STORY_MODEL, TTS_MODEL, TRANSLATION_MODEL,
+    KANNADA_TTS_MODEL, BASE_SEED, FLUX_STEPS, FLUX_GUIDANCE, FLUX_SIZE,
+    COLOR_ART_STYLE, COLOR_PAGE_SUFFIX, STORY_LENGTHS, GENRES, MOODS,
+)
+from book_builder import build_book_html, export_pdf, magic_loader_html
+from ui.layout import create_layout
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+ON_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))
+
+_FLUX       = None
+_STORY_M    = None; _STORY_TOK = None
+_TTS_M      = None
+_TRANS_M    = None; _TRANS_TOK = None
+_KAN_TTS_M  = None
+_LOAD_ERRORS = {}
+
+
+def _load_flux():
+    global _FLUX
+    if _FLUX is None:
+        from diffusers import Flux2KleinPipeline
+        _FLUX = Flux2KleinPipeline.from_pretrained(
+            FLUX_MODEL.hub_id, torch_dtype=torch.bfloat16).to("cuda")
+    return _FLUX
+
+
+def _load_story():
+    global _STORY_M, _STORY_TOK
+    if _STORY_M is None:
+        from transformers import AutoTokenizer, AutoModelForCausalLM
+        _STORY_TOK = AutoTokenizer.from_pretrained(
+            STORY_MODEL.hub_id, trust_remote_code=True)
+        _STORY_M = AutoModelForCausalLM.from_pretrained(
+            STORY_MODEL.hub_id, torch_dtype=torch.float16, trust_remote_code=True,
+        ).to("cuda").eval()
+    return _STORY_M, _STORY_TOK
+
+
+def _load_tts():
+    global _TTS_M
+    if _TTS_M is None:
+        from voxcpm import VoxCPM
+        _TTS_M = VoxCPM.from_pretrained(TTS_MODEL.hub_id, load_denoiser=False)
+    return _TTS_M
+
+
+def _load_translation():
+    global _TRANS_M, _TRANS_TOK
+    if _TRANS_M is None:
+        from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+        _TRANS_TOK = AutoTokenizer.from_pretrained(
+            TRANSLATION_MODEL.hub_id, trust_remote_code=True)
+        _TRANS_M = AutoModelForSeq2SeqLM.from_pretrained(
+            TRANSLATION_MODEL.hub_id, trust_remote_code=True,
+        ).to("cuda").eval()
+    return _TRANS_M, _TRANS_TOK
+
+
+def _load_kannada_tts():
+    global _KAN_TTS_M
+    if _KAN_TTS_M is None:
+        from indic_tts import _get_model
+        _KAN_TTS_M = _get_model()
+    return _KAN_TTS_M
+
+
+if ON_ZEROGPU:
+    for _n, _fn in [("flux", _load_flux), ("story", _load_story),
+                    ("tts", _load_tts), ("translation", _load_translation),
+                    ("kannada_tts", _load_kannada_tts)]:
+        try:
+            _fn()
+        except Exception as e:
+            _LOAD_ERRORS[_n] = repr(e)
+            logger.exception(f"Module-level load failed: {_n}")
+
+
+# ── Genre/mood → deterministic story templates (fallback) ──────────────
+
+_TEMPLATES = {
+    "Animals": [
+        ("{hero} loved exploring the meadow every evening.", "{hero} walking through a golden meadow at dusk"),
+        ("One night, {hero} heard a tiny sound in the tall grass.", "{hero} listening carefully near the rustling grass"),
+        ("A little firefly needed help finding its family.", "{hero} meeting a tiny glowing firefly"),
+        ("{hero} gently carried the firefly through the dark forest.", "{hero} walking carefully through a moonlit forest"),
+        ("Together they found the firefly's home, glowing warm and bright.", "{hero} and the firefly reunited with the glowing firefly family"),
+        ("Tired and happy, {hero} curled up under the stars.", "{hero} sleeping peacefully under a starry sky"),
+    ],
+    "Kingdom": [
+        ("In a cozy kingdom, {hero} was the kindest helper of all.", "{hero} standing cheerfully in a small fairy-tale kingdom"),
+        ("One sleepy evening, the king's golden crown went missing.", "{hero} seeing the worried king without his crown"),
+        ("{hero} searched the royal garden by moonlight.", "{hero} searching carefully through a moonlit garden"),
+        ("A sleepy mouse had borrowed it for a bed!", "{hero} discovering a tiny mouse asleep inside the crown"),
+        ("{hero} found the mouse a proper bed made of petals.", "{hero} tucking the tiny mouse into a flower-petal bed"),
+        ("The king smiled, and the whole kingdom slept in peace.", "{hero} and the king smiling together under the night sky"),
+    ],
+    "Space": [
+        ("{hero} loved watching the stars from the garden.", "{hero} lying in the grass gazing at the starry sky"),
+        ("One night, a small star blinked and fell from the sky.", "{hero} seeing a little star tumbling down"),
+        ("{hero} caught the star in a jar of moonlight.", "{hero} gently catching a glowing star in a jar"),
+        ("The star was lost and didn't know how to get home.", "{hero} listening to the sad little star"),
+        ("{hero} climbed the tallest hill and let the star go free.", "{hero} releasing the star from the hilltop into the sky"),
+        ("The star zoomed home, and {hero} fell fast asleep smiling.", "{hero} smiling and drifting off to sleep under the stars"),
+    ],
+    "Dragons": [
+        ("{hero} lived near a mountain where a shy dragon slept.", "{hero} looking up at a misty mountain at night"),
+        ("One evening, the dragon sneezed and lost its flame.", "{hero} watching the dragon sneeze sadly"),
+        ("{hero} brought warm soup and a soft blanket to the dragon.", "{hero} carrying a steaming bowl of soup to the dragon"),
+        ("The dragon felt better and puffed a tiny grateful flame.", "{hero} and the dragon sharing a warm cozy moment"),
+        ("Together they lit the lanterns along the sleepy village path.", "{hero} and the dragon lighting lanterns in the quiet village"),
+        ("The dragon curled up, and {hero} tucked it in with a smile.", "{hero} tucking the dragon in for the night"),
+    ],
+    "Ocean": [
+        ("{hero} sat by the shore watching the moonlight on the waves.", "{hero} sitting peacefully by the ocean at night"),
+        ("A little fish splashed up and looked worried.", "{hero} seeing a small worried fish near the surface"),
+        ("The fish had lost its shell-home in a big wave.", "{hero} listening to the little fish explain its problem"),
+        ("{hero} dove gently and found the shell on the sandy floor.", "{hero} swimming carefully along the moonlit ocean floor"),
+        ("The fish swam home, and the sea became calm and quiet.", "{hero} watching the happy fish return to its shell"),
+        ("{hero} fell asleep to the soft sound of the waves.", "{hero} sleeping peacefully beside the calm, moonlit sea"),
+    ],
+    "Forest": [
+        ("{hero} walked into the whispering forest as the moon rose.", "{hero} stepping into a moonlit forest path"),
+        ("The trees were worried — an owl had lost its song.", "{hero} hearing the trees whisper about the silent owl"),
+        ("{hero} climbed a mossy rock and hummed a gentle tune.", "{hero} humming softly on a mossy rock under the moon"),
+        ("The owl listened and slowly remembered its melody.", "{hero} watching the owl open its eyes and begin to sing"),
+        ("The whole forest filled with soft nighttime music.", "{hero} smiling as the forest glows with peaceful sound"),
+        ("{hero} yawned and drifted off to sleep among the roots.", "{hero} sleeping curled up peacefully at the base of a great tree"),
+    ],
+}
+
+FEW_SHOT = """
+Write a 6-page children's bedtime story for age 5 about Luna the cat. Genre: Animals. Mood: Calming.
+
+Return ONLY valid JSON:
+{
+  "title": "Luna and the Sleepy Firefly",
+  "character_description": "A small grey cat named Luna with soft fur, big green eyes, and a white tip on her tail",
+  "pages": [
+    {"page": 1, "text": "Luna loved sitting in the garden when the moon came out.", "scene": "Luna sitting in a moonlit garden"},
+    {"page": 2, "text": "One night, she heard a tiny buzzing sound in the flowers.", "scene": "Luna listening near a flower patch"},
+    {"page": 3, "text": "A little firefly was lost and couldn't find its family.", "scene": "Luna meeting a tiny glowing firefly"},
+    {"page": 4, "text": "Luna walked gently through the dark, lighting the way.", "scene": "Luna walking with the firefly glowing beside her"},
+    {"page": 5, "text": "They found the firefly's home, glowing warm and bright.", "scene": "Luna and firefly arriving at a cluster of glowing lights"},
+    {"page": 6, "text": "Luna purred softly and curled up under the stars.", "scene": "Luna sleeping peacefully under a starry sky"}
+  ]
+}
+"""
+
+
+def _build_story_locally(hero_name: str, genre: str) -> dict:
+    hero = (hero_name or "Little Hero").strip() or "Little Hero"
+    beats = _TEMPLATES.get(genre, _TEMPLATES["Animals"])
+    pages = [
+        {"page": i+1, "text": t.format(hero=hero), "scene": s.format(hero=hero)}
+        for i, (t, s) in enumerate(beats)
+    ]
+    return {
+        "title": f"{hero}'s Bedtime Dream",
+        "character_description": (
+            f"{hero}, a friendly children's storybook hero with bright colors, "
+            "bold outlines, and a cheerful expressive face"
+        ),
+        "pages": pages,
+    }
+
+
+def _parse_story_json(raw: str) -> dict | None:
+    import re
+    m = re.search(r'\{[\s\S]*\}', raw or "")
+    if not m:
+        return None
+    try:
+        d = json.loads(m.group(0))
+        if "pages" in d and "title" in d:
+            return d
+    except Exception:
+        pass
+    return None
+
+
+# ── ZeroGPU inference ──────────────────────────────────────────────────
+
+@spaces.GPU(duration=60)
+def _gen_story_gpu(hero_name: str, genre: str, mood: str) -> dict:
+    try:
+        model, tok = _load_story()
+        prompt = (
+            f"{FEW_SHOT}\n\n"
+            f"Write a 6-page children's bedtime story for age 5 about {hero_name}. "
+            f"Genre: {genre}. Mood: {mood}. Keep it gentle and sleepy.\n\n"
+            f"Return ONLY valid JSON:\n"
+        )
+        inputs = tok.apply_chat_template(
+            [{"role": "user", "content": prompt}],
+            add_generation_prompt=True, enable_thinking=False,
+            return_dict=True, return_tensors="pt",
+        ).to("cuda")
+        with torch.no_grad():
+            out = model.generate(**inputs, max_new_tokens=800, do_sample=False)
+        response = tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+        parsed = _parse_story_json(response)
+        if parsed:
+            return parsed
+    except Exception as e:
+        logger.warning(f"Story GPU failed: {e}")
+    return _build_story_locally(hero_name, genre)
+
+
+@spaces.GPU(duration=150)
+def _gen_images_gpu(char_desc: str, scenes: list,
+                    doodle_bytes: bytes | None, seed: int) -> list:
+    import io
+    from PIL import Image
+    pipe = _load_flux()
+    canonical = None
+    if doodle_bytes:
+        try:
+            ref = Image.open(io.BytesIO(doodle_bytes)).convert("RGB")
+            canonical = pipe(
+                prompt=(
+                    "Turn this child's drawing into a clean, full-body cartoon character "
+                    "for a children's storybook. Keep the EXACT same creature. "
+                    f"{COLOR_ART_STYLE}, plain white background, full character visible, centered."
+                ),
+                image=ref, height=FLUX_SIZE, width=FLUX_SIZE,
+                guidance_scale=FLUX_GUIDANCE, num_inference_steps=FLUX_STEPS,
+                generator=torch.Generator("cuda").manual_seed(seed),
+            ).images[0]
+        except Exception as e:
+            logger.warning(f"Canonical pass failed ({e}); text2img fallback")
+    images = []
+    for i, scene in enumerate(scenes):
+        if canonical is not None:
+            kw = dict(image=canonical,
+                      prompt=f"The same character. {scene}. {COLOR_ART_STYLE}, {COLOR_PAGE_SUFFIX}")
+        else:
+            kw = dict(prompt=f"{char_desc}. Scene: {scene}. {COLOR_ART_STYLE}, centered.")
+        kw.update(height=FLUX_SIZE, width=FLUX_SIZE, guidance_scale=FLUX_GUIDANCE,
+                  num_inference_steps=FLUX_STEPS,
+                  generator=torch.Generator("cuda").manual_seed(seed + i + 1))
+        images.append(pipe(**kw).images[0])
+        logger.info(f"Page {i+1}/{len(scenes)} illustrated")
+    return images
+
+
+@spaces.GPU(duration=120)
+def _gen_tts_gpu(text: str, ref_wav: str | None,
+                 mood: str, energy: float, language: str) -> str:
+    if language == "Kannada":
+        from indic_text import translate_to_kannada
+        from indic_tts import narrate_kannada
+        kannada_text = translate_to_kannada(text)
+        ref_txt = "ಇದು ನನ್ನ ಧ್ವನಿ"
+        return narrate_kannada(ref_wav or "", ref_txt, kannada_text, mood, energy)
+    else:
+        from tts import clone_and_speak
+        return clone_and_speak(
+            ref_wav=ref_wav,
+            text=text,
+            speed=0.9,
+            mood=mood.lower(),
+            energy=energy,
+        )
+
+
+# ── heartbeat helper ───────────────────────────────────────────────────
+
+def _with_heartbeat(blocking_fn, frame_fn, poll=4.0):
+    box = {}
+    def _run():
+        try:   box["val"] = blocking_fn()
+        except BaseException as e: box["err"] = e
+    th = threading.Thread(target=_run, daemon=True)
+    th.start()
+    t0 = time.time()
+    while th.is_alive():
+        th.join(timeout=poll)
+        if th.is_alive():
+            yield ("hb", frame_fn(int(time.time() - t0)))
+    if "err" in box:
+        raise box["err"]
+    yield ("done", box["val"])
+
+
+# ── main generator ─────────────────────────────────────────────────────
+
+def create_book(doodle_image, ref_audio, hero_name, genre, mood, language, length_label):
+    t0 = time.perf_counter()
+    hero_name = (hero_name or "").strip() or "Little Hero"
+    energy    = 0.45
+
+    trace = {
+        "backend": "zerogpu", "hero": hero_name,
+        "genre": genre, "mood": mood, "language": language,
+        "seed": BASE_SEED, "ts": time.strftime("%Y-%m-%d %H:%M:%S"),
+    }
+    if _LOAD_ERRORS:
+        trace["load_errors"] = _LOAD_ERRORS
+
+    _no   = gr.update(visible=False)
+    _keep = gr.update()
+
+    yield (magic_loader_html("story", hero_name),
+           "Writing the bedtime story…", None, _no, {}, "")
+
+    try:
+        story = _gen_story_gpu(hero_name, genre, mood)
+    except Exception as e:
+        yield (f"<div class='page-loading'>Error: {e}</div>",
+               f"Error: {e}", None, _no, {}, "")
+        return
+
+    title      = story.get("title", "A Bedtime Story")
+    pages      = story.get("pages", [])
+    char_desc  = story.get("character_description", "")
+    scenes     = [p.get("scene", "") for p in pages]
+    page_texts = [p.get("text", "")  for p in pages]
+    full_text  = f"{title}. {' '.join(page_texts)}"
+    trace.update(title=title, char_desc=char_desc)
+
+    yield (magic_loader_html("images", hero_name),
+           f"{title} — illustrating…", None, _no, story, json.dumps(trace, indent=2))
+
+    doodle_bytes = None
+    if doodle_image is not None:
+        import io
+        from PIL import Image
+        buf = io.BytesIO()
+        Image.fromarray(doodle_image).save(buf, format="PNG")
+        doodle_bytes = buf.getvalue()
+
+    # narration starts in parallel with illustration
+    voice_box = {}
+    def _do_voice():
+        try:
+            voice_box["path"] = _gen_tts_gpu(full_text, ref_audio, mood, energy, language)
+        except Exception as e:
+            voice_box["err"] = e
+
+    voice_th = threading.Thread(target=_do_voice, daemon=True)
+    voice_th.start()
+
+    def _audio_now():
+        return voice_box.get("path")
+
+    img_bytes, engine = None, "sketch"
+    try:
+        for kind, payload in _with_heartbeat(
+            lambda: _gen_images_gpu(char_desc, scenes, doodle_bytes, BASE_SEED),
+            lambda s: (
+                magic_loader_html("images", hero_name),
+                f"{title} — illustrating… {s}s"
+                + ("  · narration ready ▶" if _audio_now() else "  · recording…"),
+                _audio_now(), _no, story, json.dumps(trace, indent=2),
+            ),
+        ):
+            if kind == "hb":
+                yield payload
+            else:
+                import io
+                img_bytes = []
+                for img in payload:
+                    buf = io.BytesIO(); img.save(buf, format="PNG")
+                    img_bytes.append(buf.getvalue())
+                engine = "flux"
+    except Exception as e:
+        logger.exception("Image generation failed")
+        trace["image_error"] = repr(e)
+        from services.images import generate_placeholder_images
+        img_bytes = generate_placeholder_images(char_desc, scenes, doodle_bytes)
+
+    book_html = build_book_html(img_bytes, page_texts, title, engine)
+
+    while voice_th.is_alive():
+        voice_th.join(timeout=4)
+        if voice_th.is_alive():
+            yield (book_html, f"{title} — finishing narration…",
+                   _audio_now(), _no, story, json.dumps(trace, indent=2))
+    audio_path = _audio_now()
+    if voice_box.get("err"):
+        trace["tts_error"] = repr(voice_box["err"])
+
+    pdf_path = None
+    try:
+        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
+            pdf_path = export_pdf(img_bytes, page_texts, title, f.name)
+    except Exception as e:
+        logger.warning(f"PDF failed: {e}")
+
+    trace["total_sec"] = round(time.perf_counter() - t0, 2)
+    trace["engine"] = engine
+
+    pdf_update = gr.update(value=pdf_path, visible=True) if pdf_path else _keep
+
+    yield (
+        book_html,
+        f"Done: {title} · {len(img_bytes)} pages · {language} · {trace['total_sec']}s",
+        audio_path, pdf_update, story, json.dumps(trace, indent=2),
+    )
+
+
+if __name__ == "__main__":
+    demo = create_layout(create_book_fn=create_book)
+    demo.queue(default_concurrency_limit=2, max_size=8)
+    demo.launch(share=False, allowed_paths=[tempfile.gettempdir()],
+                **demo.design_kwargs)