Update app.py

app.py

@@ -3,17 +3,9 @@ import torch
 import gradio as gr
 import datetime
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
-from safetensors.torch import load_file
 
 import spaces
 
-@spaces.GPU
-def use_gpu():
-    import torch
-    print("Torch CUDA available:", torch.cuda.is_available())
-    return {"cuda_available": torch.cuda.is_available()}
-
-
 # Constants
 MODEL_CONFIG = {
     "G0-Release": "FlameF0X/Snowflake-G0-Release",
@@ -45,48 +37,54 @@ css = """
 .model-select { background-color: #2a2a4a; padding: 10px; border-radius: 8px; margin-bottom: 15px; }
 """
 
+# Global registry - models will be loaded on-demand within GPU function
 model_registry = {}
 
-def load_all_models():
-    for name, model_id in MODEL_CONFIG.items():
-        print(f"Loading model: {name} from {model_id}")
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        safetensor_path = os.path.join(model_id, "model.safetensors")
-        if os.path.exists(safetensor_path):
-            print("Loading from safetensors...")
-            model = load_file(safetensor_path)
-        else:
-            print("Loading from Hugging Face or .bin...")
-            # Key fix: no device_map, load on CPU only
-            model = AutoModelForCausalLM.from_pretrained(
-                model_id,
-                torch_dtype=torch.float32,
-                device_map=None
-            )
+def load_model_cpu(model_id):
+    """Load model on CPU only - no CUDA initialization"""
+    print(f"Loading model on CPU: {model_id}")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Load model on CPU only
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.float32,
+        device_map=None,  # No device mapping
+        low_cpu_mem_usage=True
+    )
+    
+    return model, tokenizer
 
+@spaces.GPU
+def generate_text_gpu(prompt, model_version, temperature, top_p, top_k, max_new_tokens):
+    """GPU-decorated function for text generation"""
+    try:
+        # Load model if not already loaded
+        if model_version not in model_registry:
+            model_id = MODEL_CONFIG[model_version]
+            model, tokenizer = load_model_cpu(model_id)
+            model_registry[model_version] = (model, tokenizer)
+        
+        model, tokenizer = model_registry[model_version]
+        
+        # Move model to GPU only inside this function
+        if torch.cuda.is_available():
+            model = model.cuda()
+            device = "cuda"
+        else:
+            device = "cpu"
+        
+        # Create pipeline inside GPU function
         pipeline = TextGenerationPipeline(
             model=model,
             tokenizer=tokenizer,
             return_full_text=False,
-            max_length=MAX_LENGTH
+            max_length=MAX_LENGTH,
+            device=device
         )
-
-        model_registry[name] = (model, tokenizer, pipeline)
-
-def generate_text(prompt, model_version, temperature, top_p, top_k, max_new_tokens, history=None):
-    if history is None:
-        history = []
-    history.append({"role": "user", "content": prompt})
-
-    try:
-        if model_version not in model_registry:
-            raise ValueError(f"Model '{model_version}' not found.")
-
-        _, tokenizer, pipeline = model_registry[model_version]
-
+        
         outputs = pipeline(
             prompt,
             do_sample=temperature > 0,
@@ -97,19 +95,43 @@ def generate_text(prompt, model_version, temperature, top_p, top_k, max_new_toke
             pad_token_id=tokenizer.pad_token_id,
             num_return_sequences=1
         )
-
+        
         response = outputs[0]["generated_text"]
-        history.append({"role": "assistant", "content": response, "model": model_version})
+        return response, None
+        
+    except Exception as e:
+        error_msg = f"Error generating response: {str(e)}"
+        return error_msg, str(e)
 
+def generate_text(prompt, model_version, temperature, top_p, top_k, max_new_tokens, history=None):
+    """Main generation function that calls GPU function"""
+    if history is None:
+        history = []
+    
+    # Add user message to history
+    history.append({"role": "user", "content": prompt})
+    
+    try:
+        # Call GPU function
+        response, error = generate_text_gpu(
+            prompt, model_version, temperature, top_p, top_k, max_new_tokens
+        )
+        
+        if error:
+            history.append({"role": "assistant", "content": f"[ERROR] {response}", "model": model_version})
+        else:
+            history.append({"role": "assistant", "content": response, "model": model_version})
+        
+        # Format history for display
         formatted_history = []
         for entry in history:
             prefix = "👤 User: " if entry["role"] == "user" else f"❄️ [{entry.get('model', 'Model')}]: "
             formatted_history.append(f"{prefix}{entry['content']}")
-
+        
         return response, history, "\n\n".join(formatted_history)
-
+        
     except Exception as e:
-        error_msg = f"Error generating response: {str(e)}"
+        error_msg = f"Error in generation pipeline: {str(e)}"
         history.append({"role": "assistant", "content": f"[ERROR] {error_msg}", "model": model_version})
         return error_msg, history, str(history)
 
@@ -230,21 +252,9 @@ def create_demo():
 
     return demo
 
-# Initialize
-print("Loading Snowflake models...")
-try:
-    load_all_models()
-    print("All models loaded successfully!")
-    demo = create_demo()
-except Exception as e:
-    print(f"Error loading models: {e}")
-    with gr.Blocks(css=css) as demo:
-        gr.HTML(f"""
-        <div class="header" style="background-color: #ffebee;">
-            <h1><span class="snowflake-icon">⚠️</span> Error Loading Models</h1>
-            <p>There was a problem loading the Snowflake models: {str(e)}</p>
-        </div>
-        """)
+# Initialize demo without loading models (they'll load on-demand)
+print("Initializing Snowflake Models Demo...")
+demo = create_demo()
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()