Fixed import issues

Author: UtkarshTheDev

locallab/config.py

@@ -208,7 +208,8 @@ def estimate_model_requirements(model_id: str) -> Dict[str, float]:
         "ram_gb": requirements["min_ram"] * 1.2,  # 20% buffer
         "vram_gb": requirements["min_vram"] * 1.2 if "min_vram" in requirements else 0
     }
-
+import torch
+import psutil
 
 # Model Configuration
 CUSTOM_MODEL = get_env_var("LOCALLAB_CUSTOM_MODEL", default="")
@@ -600,7 +601,7 @@ def set_env_var(name: str, value: str):
 def get_hf_token(interactive: bool = False) -> Optional[str]:
     """Get HuggingFace token from environment or config"""
     # First check environment
-    token = get_env_var(HF_TOKEN_ENV)
+    token = os.environ.get("HUGGINGFACE_TOKEN", "").strip()
 
     # Then check config
     if not token:
@@ -609,16 +610,45 @@ def get_hf_token(interactive: bool = False) -> Optional[str]:
             token = str(get_config_value("huggingface_token", "")).strip()
             if token:
                 # Update environment variable
-                set_env_var(HF_TOKEN_ENV, token)
+                os.environ["HUGGINGFACE_TOKEN"] = token
         except:
             pass
 
+    # If interactive and still no token, prompt user
+    if not token and interactive:
+        try:
+            click.echo("\n🔑 HuggingFace token is required for accessing this model.")
+            click.echo("Get your token from: https://huggingface.co/settings/tokens")
+            
+            token = click.prompt(
+                "Enter your HuggingFace token",
+                type=str,
+                default="",
+                show_default=False
+            ).strip()
+            
+            if token:
+                if len(token) < 20:
+                    click.echo("\n❌ Invalid token format. Please check your token.")
+                    return None
+                    
+                click.echo(f"\n✅ Token saved: {token}")
+                os.environ["HUGGINGFACE_TOKEN"] = token
+                
+                # Save to config
+                from .cli.config import set_config_value
+                set_config_value("huggingface_token", token)
+            else:
+                click.echo("\nSkipping token...")
+        except:
+            pass
+            
     return token
 
 def get_ngrok_token() -> Optional[str]:
     """Get ngrok token from environment or config"""
     # First check environment
-    token = get_env_var(NGROK_TOKEN_ENV)
+    token = os.environ.get("NGROK_AUTHTOKEN", "").strip()
 
     # Then check config
     if not token:
@@ -627,14 +657,16 @@ def get_ngrok_token() -> Optional[str]:
             token = str(get_config_value("ngrok_auth_token", "")).strip()
             if token:
                 # Update environment variable
-                set_env_var(NGROK_TOKEN_ENV, token)
+                os.environ["NGROK_AUTHTOKEN"] = token
         except:
             pass
 
     return token
 
 def save_config(config: Dict[str, Any]):
     """Save configuration to file"""
+    from .cli.config import ensure_config_dir, CONFIG_FILE
+    
     ensure_config_dir()
 
     # Ensure tokens are stored as proper strings
@@ -653,51 +685,3 @@ def save_config(config: Dict[str, Any]):
             json.dump(config, f, indent=2)
     except Exception as e:
         logger.error(f"Error saving config: {e}")
-
-
-def get_hf_token(interactive: bool = False) -> Optional[str]:
-    """Get HuggingFace token from environment or config"""
-    # First check environment
-    token = os.environ.get("HUGGINGFACE_TOKEN", "").strip()
-    
-    # Then check config
-    if not token:
-        try:
-            from .cli.config import get_config_value
-            token = str(get_config_value("huggingface_token", "")).strip()
-            if token:
-                # Update environment variable
-                os.environ["HUGGINGFACE_TOKEN"] = token
-        except:
-            pass
-    
-    # If interactive and still no token, prompt user
-    if not token and interactive:
-        try:
-            click.echo("\n🔑 HuggingFace token is required for accessing this model.")
-            click.echo("Get your token from: https://huggingface.co/settings/tokens")
-            
-            token = click.prompt(
-                "Enter your HuggingFace token",
-                type=str,
-                default="",
-                show_default=False
-            ).strip()
-            
-            if token:
-                if len(token) < 20:
-                    click.echo("\n❌ Invalid token format. Please check your token.")
-                    return None
-                    
-                click.echo(f"\n✅ Token saved: {token}")
-                os.environ["HUGGINGFACE_TOKEN"] = token
-                
-                # Save to config
-                from .cli.config import set_config_value
-                set_config_value("huggingface_token", token)
-            else:
-                click.echo("\nSkipping token...")
-        except:
-            pass
-            
-    return token

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="locallab",
-    version="0.4.32",
+    version="0.4.33",
     packages=find_packages(include=["locallab", "locallab.*"]),
     install_requires=[
         "fastapi>=0.95.0,<1.0.0",