change script

app.py

@@ -17,6 +17,16 @@ from scripts.generate_subtitles import (
     add_subtitles_to_video,
 )
 
+from __future__ import annotations
+import os
+from typing import List
+import torch
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    BitsAndBytesConfig,
+)
+
 # ──────────────────────────────────────────────────────────────────────────────
 # Constants & utilities
 # ──────────────────────────────────────────────────────────────────────────────
@@ -28,6 +38,30 @@ ASSETS_DIRS = (
     "./assets/video_music",
 )
 
+
+# ────────────────────────────────────────────────────────
+#   CONFIGURATION
+# ────────────────────────────────────────────────────────
+MODEL_ID  = os.getenv("MODEL_ID", "unsloth/Qwen3-30B-A3B")
+USE_INT8  = os.getenv("USE_INT8", "0") == "1"    # set env var to 1 for 8‑bit
+THINKING  = os.getenv("DEFAULT_THINKING", "1") == "1"  # default reasoning mode
+
+DTYPE     = torch.bfloat16  # Hopper bf16 fast‑path
+bnb_cfg   = BitsAndBytesConfig(load_in_8bit=True) if USE_INT8 else None
+
+print(f"🔄 Loading {MODEL_ID} ({'8‑bit' if USE_INT8 else 'bf16'}) …")
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    device_map="auto",
+    torch_dtype=(None if USE_INT8 else DTYPE),
+    quantization_config=bnb_cfg,
+    trust_remote_code=True,
+).eval()
+DEVICE = model.device
+print("✅ Model ready.")
+
 for d in ASSETS_DIRS:
     os.makedirs(d, exist_ok=True)
 
@@ -69,10 +103,10 @@ def cb_generate_script(
             f"Instruction: {instruction.strip()}\n"
             f"🔴 Strict target duration: {target_duration}s — ≈ {approx_words} words (must be respected)."
         )
-        script = generate_script(prompt)
+        script = generate_script(model,tokenizer,prompt)
 
-    title = generate_title(script)
-    description = generate_description(script)
+    title = generate_title(model,tokenizer,script)
+    description = generate_description(model,tokenizer,script)
     return script, title, description, script  # last return for state update
 
 

scripts/generate_scripts.py

@@ -12,14 +12,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
 @spaces.GPU()
-def generate_local(prompt: str, max_new_tokens: int = 350, temperature: float = 0.7) -> str:
-    model_id = "Qwen/Qwen3-0.6B"
+def generate_local(model,tokenizer, prompt: str, max_new_tokens: int = 350, temperature: float = 0.7) -> str:
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # get the device the model is on
-
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, trust_remote_code=True).to(device)
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    
     output_ids = model.generate(
         **inputs,
         max_new_tokens=max_new_tokens,
@@ -31,7 +26,7 @@ def generate_local(prompt: str, max_new_tokens: int = 350, temperature: float =
 
 
 
-def generate_script(prompt: str, word_count: int = 60) -> str:
+def generate_script(model,tokenizer, prompt: str, word_count: int = 60) -> str:
     system_prompt = (
         "You are an expert YouTube scriptwriter. "
         "Your job is to write the EXACT words that will be spoken aloud in a video. "
@@ -44,15 +39,15 @@ def generate_script(prompt: str, word_count: int = 60) -> str:
         "- Do NOT include any explanations, labels, or headers. Only output the final spoken script.\n\n"
         "Start now:"
     )
-    return generate_local(system_prompt)
+    return generate_local(model,tokenizer, system_prompt)
 
 
-def one_word(query: str) -> str:
+def one_word(model,tokenizer, query: str) -> str:
     prompt_final = (
         "Extract only the unique central theme of the following text in English in JSON format like this: "
         '{"keyword": "impact"}. Text: ' + query
     )
-    result = generate_local(prompt_final, max_new_tokens=30, temperature=0.4)
+    result = generate_local(model,tokenizer, prompt_final, max_new_tokens=30, temperature=0.4)
     try:
         keyword_json = json.loads(result)
         keyword = keyword_json.get("keyword", "")
@@ -62,14 +57,14 @@ def one_word(query: str) -> str:
     return keyword.lower()
 
 
-def generate_title(text: str) -> str:
+def generate_title(model,tokenizer, text: str) -> str:
     prompt_final = (
         "Generate a unique title for a YouTube Short video that is engaging and informative, "
         "maximum 100 characters, without emojis, introduction, or explanation. Content:\n" + text
     )
-    return generate_local(prompt_final, max_new_tokens=50, temperature=0.9).strip()
+    return generate_local(model,tokenizer, prompt_final, max_new_tokens=50, temperature=0.9).strip()
 
-def generate_description(text: str) -> str:
+def generate_description(model,tokenizer, text: str) -> str:
     prompt_final = (
         "Write only the YouTube video description in English:\n"
         "1. A compelling opening line.\n"
@@ -77,13 +72,13 @@ def generate_description(text: str) -> str:
         "3. End with 3 relevant hashtags.\n"
         "No emojis or introductions. Here is the text:\n" + text
     )
-    return generate_local(prompt_final, max_new_tokens=300, temperature=0.7).strip()
+    return generate_local(model,tokenizer, prompt_final, max_new_tokens=300, temperature=0.7).strip()
 
-def generate_tags(text: str) -> list:
+def generate_tags(model,tokenizer, text: str) -> list:
     prompt_final = (
         "List only the important keywords for this YouTube video, separated by commas, "
         "maximum 10 keywords. Context: " + text
     )
-    result = generate_local(prompt_final, max_new_tokens=100, temperature=0.5)
+    result = generate_local(model,tokenizer, prompt_final, max_new_tokens=100, temperature=0.5)
     return [tag.strip() for tag in result.split(",") if tag.strip()]