Update app.py

app.py

@@ -1,79 +1,72 @@
 import gradio as gr
 import torch
+import torch.multiprocessing as mp
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
 import logging
+import os
+
+# Set multiprocessing to 'spawn' for ZeroGPU compatibility
+mp.set_start_method('spawn', force=True)
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Global variables for model and tokenizer (lazy loading)
+# Global variables for model and tokenizer (load at startup)
 model = None
 tokenizer = None
 MODEL_NAME = "ubiodee/Test_Plutus"
 FALLBACK_TOKENIZER = "NousResearch/Meta-Llama-3-8B"
 
-# Load tokenizer at startup (lightweight, no model yet)
+# Load tokenizer at startup
 try:
     logger.info("Loading tokenizer at startup for %s...", MODEL_NAME)
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
-        use_fast=True,  # Llama-3 uses fast tokenizer
+        use_fast=True,
         trust_remote_code=True,
     )
     logger.info("Primary tokenizer loaded successfully.")
 except Exception as e:
     logger.warning(f"Primary tokenizer failed: {str(e)}. Using fallback: {FALLBACK_TOKENIZER}")
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(
-            FALLBACK_TOKENIZER,
-            use_fast=True,
-            trust_remote_code=True,
-        )
-        logger.info("Fallback tokenizer loaded successfully.")
-    except Exception as fallback_e:
-        logger.error(f"Fallback tokenizer failed: {str(fallback_e)}")
-        raise
+    tokenizer = AutoTokenizer.from_pretrained(
+        FALLBACK_TOKENIZER,
+        use_fast=True,
+        trust_remote_code=True,
+    )
+    logger.info("Fallback tokenizer loaded successfully.")
 
 # Set pad token
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
     logger.info("Set pad_token_id to eos_token_id: %s", tokenizer.eos_token_id)
 
-def load_model():
-    """Load model inside GPU context."""
-    global model
-    if model is None:
-        try:
-            logger.info("Loading model %s with torch.float16...", MODEL_NAME)
-            model = AutoModelForCausalLM.from_pretrained(
-                MODEL_NAME,
-                torch_dtype=torch.float16,  # Use fp16 for ZeroGPU
-                low_cpu_mem_usage=True,
-                trust_remote_code=True,
-            )
-            model.eval()
-            if torch.cuda.is_available():
-                model.to("cuda")
-                logger.info("Model loaded and moved to GPU.")
-            else:
-                logger.warning("GPU not available; using CPU.")
-        except Exception as e:
-            logger.error(f"Model loading failed: {str(e)}")
-            raise
-    return model
+# Load model at startup (CPU/fp16, move to GPU in decorated function)
+try:
+    logger.info("Loading model %s with torch.float16 on CPU...", MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+        trust_remote_code=True,
+        device_map="cpu",  # Load on CPU to avoid CUDA init issues
+    )
+    model.eval()
+    logger.info("Model loaded successfully on CPU.")
+except Exception as e:
+    logger.error(f"Model loading failed: {str(e)}")
+    raise
 
-# Response function: Load model on first call, then reuse
-@spaces.GPU(duration=300)  # Allow up to 5min for loading + inference
+# Response function: Transfer to GPU and infer (no CUDA init here)
+@spaces.GPU(duration=120)  # Reduced for quota efficiency
 def generate_response(prompt, progress=gr.Progress()):
     global model
-    progress(0.1, desc="Loading model if needed...")
-    model = load_model()  # Ensures model is loaded in GPU context
-    
-    progress(0.3, desc="Tokenizing input...")
+    progress(0.1, desc="Moving model to GPU...")
     try:
-        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
+        model = model.to("cuda")  # Move to GPU in decorated context
+        progress(0.3, desc="Tokenizing input...")
+        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512).to("cuda")
         
         progress(0.6, desc="Generating response...")
         with torch.no_grad():
@@ -97,6 +90,10 @@ def generate_response(prompt, progress=gr.Progress()):
     except Exception as e:
         logger.error(f"Inference failed: {str(e)}")
         return f"Error during generation: {str(e)}"
+    finally:
+        # Clean up GPU memory
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
 # Gradio UI
 demo = gr.Interface(
@@ -107,5 +104,5 @@ demo = gr.Interface(
     description="Write Plutus smart contracts on Cardano blockchain."
 )
 
-# Launch with ZeroGPU-compatible settings
+# Launch without queue args (ZeroGPU handles it)
 demo.launch()