Update app.py

app.py

@@ -2,19 +2,9 @@ import os
 import torch
 import gradio as gr
 import datetime
-from spaces import GPU
 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",
@@ -28,13 +18,11 @@ TOP_P_DEFAULT = 0.9
 TOP_K_DEFAULT = 40
 MAX_NEW_TOKENS_DEFAULT = 256
 
-# UI parameter bounds
 TEMPERATURE_MIN, TEMPERATURE_MAX = 0.1, 2.0
 TOP_P_MIN, TOP_P_MAX = 0.1, 1.0
 TOP_K_MIN, TOP_K_MAX = 1, 100
 MAX_NEW_TOKENS_MIN, MAX_NEW_TOKENS_MAX = 16, 1024
 
-# Styling
 css = """
 .gradio-container { background-color: #1e1e2f !important; color: #e0e0e0 !important; }
 .header { background-color: #2b2b3c; padding: 20px; margin-bottom: 20px; border-radius: 10px; text-align: center; }
@@ -48,7 +36,6 @@ css = """
 .model-select { background-color: #2a2a4a; padding: 10px; border-radius: 8px; margin-bottom: 15px; }
 """
 
-# Model registry
 model_registry = {}
 
 def load_all_models():
@@ -64,10 +51,11 @@ def load_all_models():
             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.float16 if torch.cuda.is_available() else torch.float32,
-                device_map="auto"
+                torch_dtype=torch.float32,
+                device_map=None
             )
 
         pipeline = TextGenerationPipeline(
@@ -83,7 +71,7 @@ def generate_text(prompt, model_version, temperature, top_p, top_k, max_new_toke
     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.")
@@ -110,7 +98,7 @@ def generate_text(prompt, model_version, temperature, top_p, top_k, max_new_toke
             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)}"
         history.append({"role": "assistant", "content": f"[ERROR] {error_msg}", "model": model_version})
@@ -121,14 +109,13 @@ def clear_conversation():
 
 def create_demo():
     with gr.Blocks(css=css) as demo:
-        # Header
         gr.HTML("""
         <div class="header">
             <h1><span class="snowflake-icon">❄️</span> Snowflake Models Demo</h1>
             <p>Experience the capabilities of the Snowflake series language models</p>
         </div>
         """)
-        
+
         with gr.Column():
             with gr.Row(elem_classes="model-select"):
                 model_version = gr.Radio(
@@ -137,21 +124,21 @@ def create_demo():
                     label="Select Model Version",
                     info="Choose which Snowflake model to use"
                 )
-            
+
             chat_history_display = gr.Textbox(
-                value="", 
-                label="Conversation History", 
-                lines=10, 
+                value="",
+                label="Conversation History",
+                lines=10,
                 max_lines=30,
                 interactive=False
             )
-            
+
             history_state = gr.State([])
-            
+
             with gr.Row():
                 with gr.Column(scale=4):
                     prompt = gr.Textbox(
-                        placeholder="Type your message here...", 
+                        placeholder="Type your message here...",
                         label="Your Input",
                         lines=2
                     )
@@ -160,14 +147,13 @@ def create_demo():
                     clear_btn = gr.Button("Clear Conversation")
 
             response_output = gr.Textbox(
-                value="", 
-                label="Model Response", 
+                value="",
+                label="Model Response",
                 lines=5,
                 max_lines=10,
                 interactive=False
             )
-        
-        # Generation Parameters
+
         with gr.Accordion("Generation Parameters", open=False):
             with gr.Column(elem_classes="parameter-section"):
                 with gr.Row():
@@ -193,8 +179,7 @@ def create_demo():
                             value=MAX_NEW_TOKENS_DEFAULT, step=8,
                             label="Maximum New Tokens"
                         )
-        
-        # Example prompts
+
         examples = [
             "Write a short story about a snowflake that comes to life.",
             "Explain the concept of artificial neural networks to a 10-year-old.",
@@ -211,14 +196,13 @@ def create_demo():
                     label="Click on an example to try it",
                     examples_per_page=5
                 )
-        
+
         gr.HTML(f"""
         <div class="footer">
             <p>Snowflake Models Demo • Created with Gradio • {datetime.datetime.now().year}</p>
         </div>
         """)
-        
-        # Interactions
+
         submit_btn.click(
             fn=generate_text,
             inputs=[prompt, model_version, temperature, top_p, top_k, max_new_tokens, history_state],
@@ -253,6 +237,5 @@ except Exception as e:
         </div>
         """)
 
-# Run app
 if __name__ == "__main__":
     demo.launch()