Update app.py

app.py

@@ -24,34 +24,41 @@ library_username = os.environ.get("OLLAMA_USERNAME").lower()
 
 
 def regenerate_pubkey(login, oauth_token: gr.OAuthToken | None):
+    try:
+        if oauth_token.token is None:
+            raise ValueError("You must be logged in to use Ollamafy")
+        hash_oauth = hashlib.sha256(unicode_string.encode(oauth_token)).hexdigest()
+        ollama_host_file = f"echo $(ss -natp | grep (cat ollama.pid) | awk '{{print $4}}') > ollama.host"
+        ollama_pid_file = f"echo $! > ollama.pid"
+        ollama_pubkey = open(f"{HOME}/.ollama/{hash_oauth}/id_ed25519.pub", "r")    
+        ollama_start = f"HOME={HOME}/{hash_oauth} ollama serve & {ollama_pid_file} & sleep 5 & {ollama_host_file}"    
+        ollama_stop = f"kill -9 'cat ${HOME}/{hash_oauth}/ollama.pid'"
+        delete_home = f"rm -Rf {HOME}/{hash_oauth}/.ollama"
+            
+        result = subprocess.run(ollama_stop, shell=True, capture_output=True)
+        print(result)
+        if result.returncode != 0:
+            raise Exception(f"Error stoppping Ollama {result.stderr}")
+        print("Ollama stopped successfully!")
+    
+        result = subprocess.run(delete_home, shell=True, capture_output=True)
+        print(result)
+        if result.returncode != 0:
+            raise Exception(f"Error removing Ollama HOME folder {result.stderr}")
+        print("Ollama HOME fodler removed successfully!")
+    
+        result = subprocess.run(ollama_start, shell=True, capture_output=True)
+        print(result)
+        if result.returncode != 0:
+            raise Exception(f"Error starting Ollama {result.stderr}")
+        print("Ollama started successfully!")
+    
+    except Exception as e:
+        return (f"Error: {e}", "error.png")
+    finally:
+        shutil.rmtree(model_name, ignore_errors=True)
+        print("Folder cleaned up successfully`!")
 
-    if oauth_token.token is None:
-        raise ValueError("You must be logged in to use Ollamafy")
-    hash_oauth = hashlib.sha256(unicode_string.encode(oauth_token)).hexdigest()
-    ollama_host_file = f"echo $(ss -natp | grep (cat ollama.pid) | awk '{{print $4}}') > ollama.host"
-    ollama_pid_file = f"echo $! > ollama.pid"
-    ollama_pubkey = open(f"{HOME}/.ollama/{hash_oauth}/id_ed25519.pub", "r")    
-    ollama_start = f"HOME={HOME}/{hash_oauth} ollama serve & {ollama_pid_file} & sleep 5 & {ollama_host_file}"    
-    ollama_stop = f"kill -9 'cat ${HOME}/{hash_oauth}/ollama.pid'"
-    delete_home = f"rm -Rf {HOME}/{hash_oauth}/.ollama"
-        
-    result = subprocess.run(ollama_stop, shell=True, capture_output=True)
-    print(result)
-    if result.returncode != 0:
-        raise Exception(f"Error stoppping Ollama {result.stderr}")
-    print("Ollama stopped successfully!")
-
-    result = subprocess.run(delete_home, shell=True, capture_output=True)
-    print(result)
-    if result.returncode != 0:
-        raise Exception(f"Error removing Ollama HOME folder {result.stderr}")
-    print("Ollama HOME fodler removed successfully!")
-
-    result = subprocess.run(ollama_start, shell=True, capture_output=True)
-    print(result)
-    if result.returncode != 0:
-        raise Exception(f"Error starting Ollama {result.stderr}")
-    print("Ollama started successfully!")
 
 def ollamafy_model(login, account, model_id, ollama_library_username , ollama_q_method, latest, maintainer, oauth_token: gr.OAuthToken | None):
     ollama_library_username: library_username | None 
@@ -193,7 +200,6 @@ with gr.Blocks(css=css) as demo:
             )
         with gr.Column(scale=1):
             generate_pubkey = gr.Button (
-                fn=regenerate_pubkey,
                 value=regenerate_pubkey,
                 min_width=250,
             )