Update app.py

app.py

@@ -42,7 +42,7 @@ tokenizer = AutoTokenizer.from_pretrained(LORA_MODEL, trust_remote_code=True)
 
 print("Loading models on CPU to safely bypass ZeroGPU startup constraints...")
 
-# 🛑 THE FIX: Temporarily blind the environment to CUDA so PEFT doesn't crash
+# 🛑 Temporarily blind the environment to CUDA so PEFT doesn't crash during startup
 _original_cuda_available = torch.cuda.is_available
 torch.cuda.is_available = lambda: False 
 
@@ -61,7 +61,7 @@ try:
     _peft_model = PeftModel.from_pretrained(
         base_model, 
         LORA_MODEL, 
-        device_map={"": "cpu"} # Extra safeguard
+        device_map={"": "cpu"}
     )
 
     print("Merging LoRA weights...")
@@ -87,7 +87,7 @@ try:
     _draft_peft = PeftModel.from_pretrained(
         _draft_base, 
         DRAFT_MODEL, 
-        device_map={"": "cpu"} # Extra safeguard
+        device_map={"": "cpu"}
     )
 
     draft_model = _draft_peft.merge_and_unload()
@@ -129,6 +129,10 @@ def generate_code(topic: str, temperature: float, max_tokens: int):
         yield "# Please enter a topic or description."
         return
 
+    # 🚀 THE FINAL FIX: Move models from CPU RAM to the ZeroGPU hardware just-in-time
+    model.to(device)
+    draft_model.to(device)
+
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user",   "content": topic.strip()},
@@ -274,16 +278,7 @@ CSS = """
 }
 """
 
-with gr.Blocks(
-    title="Anim Vid Max — AI Animation Generator",
-    css=CSS,
-    theme=gr.themes.Base(
-        primary_hue=gr.themes.colors.indigo,
-        secondary_hue=gr.themes.colors.sky,
-        neutral_hue=gr.themes.colors.slate,
-        font=gr.themes.GoogleFont("Inter"),
-    ),
-) as demo:
+with gr.Blocks(title="Anim Vid Max — AI Animation Generator") as demo:
     
     gr.HTML("""
     <div class="header-wrap">
@@ -324,7 +319,7 @@ with gr.Blocks(
                 label="💡 Examples",
             )
 
-        with gr.Column(scale=3.5, elem_classes=["col-sep"]):
+        with gr.Column(scale=3, elem_classes=["col-sep"]):
             gr.HTML('<div class="section-label">2. Manim Code</div>')
             code_output = gr.Code(
                 label="Generated Script (editable)",
@@ -364,4 +359,13 @@ with gr.Blocks(
         outputs=[output_video, status_output],
     )
 
-demo.launch()
+# Cleaned up launch arguments to prevent warnings
+demo.launch(
+    css=CSS, 
+    theme=gr.themes.Base(
+        primary_hue=gr.themes.colors.indigo,
+        secondary_hue=gr.themes.colors.sky,
+        neutral_hue=gr.themes.colors.slate,
+        font=gr.themes.GoogleFont("Inter"),
+    )
+)