Update app.py

app.py

@@ -1,42 +1,29 @@
 import os
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 import gradio as gr
 
-# Configuration
-MODEL_ID = "google/gemma-1.1-2b-it"  # Using smaller 2B version
+# Configuration - Using even smaller model
+MODEL_ID = "google/gemma-1.1-2b-it"
 HF_TOKEN = os.getenv("HF_TOKEN")
-MAX_TOKENS = 100  # Conservative limit
-OFFLOAD_FOLDER = "./offload"  # Directory for disk offloading
+MAX_TOKENS = 80  # Very conservative limit
 
 def load_model():
-    """Proper model loading with memory optimization"""
-    print("🔄 Loading model with memory optimization...")
+    """Simplified model loading that works with Spaces"""
+    print("🔄 Loading model (this may take a few minutes)...")
     
-    # Create offload directory
-    os.makedirs(OFFLOAD_FOLDER, exist_ok=True)
-    
-    # Load tokenizer
+    # Load with explicit CPU mapping
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-    
-    # Special loading for memory constraints
-    with init_empty_weights():
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            torch_dtype=torch.float32,
-            token=HF_TOKEN
-        )
-    
-    model = load_checkpoint_and_dispatch(
-        model,
-        checkpoint=MODEL_ID,
-        device_map="auto",
-        offload_folder=OFFLOAD_FOLDER,
-        no_split_module_classes=["GemmaDecoderLayer"]
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        device_map="cpu",  # Force CPU-only
+        torch_dtype=torch.float32,  # Required for CPU
+        token=HF_TOKEN
     )
     
-    print("✅ Model loaded with disk offloading!")
+    # Ensure weights are tied
+    model.tie_weights()
+    print("✅ Model loaded successfully!")
     return tokenizer, model
 
 tokenizer, model = load_model()
@@ -44,8 +31,8 @@ tokenizer, model = load_model()
 def predict(topic):
     """Memory-safe generation"""
     try:
-        prompt = f"Create a short script about {topic}:"
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        prompt = f"Create a VERY short script about {topic}:\n1) Hook\n2) Point\n3) CTA\n\nScript:"
+        inputs = tokenizer(prompt, return_tensors="pt")
         
         outputs = model.generate(
             **inputs,
@@ -59,10 +46,10 @@ def predict(topic):
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Gradio interface
+# Minimal interface
 gr.Interface(
     fn=predict,
     inputs=gr.Textbox(label="Topic"),
-    outputs=gr.Textbox(label="Script", lines=5),
+    outputs=gr.Textbox(label="Script", lines=4),
     api_name="predict"
 ).launch(server_name="0.0.0.0")