image model gated

leveraging_machine_learning.py

@@ -20,12 +20,12 @@ modelnames = ["stvlynn/Gemma-2-2b-Chinese-it", "unsloth/Llama-3.2-1B-Instruct",
               "Qwen/Qwen2.5-7B-Instruct", "Qwen/Qwen2-0.5B-Instruct", "Qwen/Qwen2-1.5B-Instruct", "Qwen/Qwen2-7B-Instruct", "Qwen/Qwen1.5-MoE-A2.7B-Chat", "HuggingFaceTB/SmolLM-135M-Instruct", "microsoft/Phi-3-mini-4k-instruct", "Groq/Llama-3-Groq-8B-Tool-Use", "hugging-quants/Meta-Llama-3.1-8B-Instruct-BNB-NF4", 
               "SpectraSuite/TriLM_3.9B_Unpacked", "h2oai/h2o-danube3-500m-chat", "OuteAI/Lite-Mistral-150M-v2-Instruct", "Zyphra/Zamba2-1.2B", "anthracite-org/magnum-v2-4b", ]
 
-imagemodelnames = ["black-forest-labs/FLUX.1-schnell"]
+# imagemodelnames = ["black-forest-labs/FLUX.1-schnell"]
 
 current_model_index = 0
 current_image_model_index = 0
 modelname = modelnames[current_model_index]
-imagemodelname = imagemodelnames[current_image_model_index]
+# imagemodelname = imagemodelnames[current_image_model_index]
 lastmodelnameinloadfunction = None
 lastimagemodelnameinloadfunction = None
 
@@ -35,7 +35,7 @@ embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
 # Initialize model and tokenizer as global variables
 model = None
 tokenizer = None
-flux_pipe = None
+# flux_pipe = None
 
 # Dictionary to store loaded models
 loaded_models = {}
@@ -90,42 +90,42 @@ def load_model(model_name):
             f"Tokenizer size: {get_size_str(tokenizer_size)}, "
             f"GPU memory used: {get_size_str(memory_used)}")
 
-def load_image_model(imagemodelname):
-    global flux_pipe, lastimagemodelnameinloadfunction, loaded_models
+# def load_image_model(imagemodelname):
+#     global flux_pipe, lastimagemodelnameinloadfunction, loaded_models
 
-    print(f"Loading image model: {imagemodelname}")
+#     print(f"Loading image model: {imagemodelname}")
     
-    # Record initial GPU memory usage
-    initial_memory = torch.cuda.memory_allocated()
+#     # Record initial GPU memory usage
+#     initial_memory = torch.cuda.memory_allocated()
 
-    if 'flux_pipe' in globals() and flux_pipe is not None:
-        flux_pipe = None
+#     if 'flux_pipe' in globals() and flux_pipe is not None:
+#         flux_pipe = None
     
-    torch.cuda.empty_cache()
-    gc.collect()
+#     torch.cuda.empty_cache()
+#     gc.collect()
 
-    flux_pipe = FluxPipeline.from_pretrained(imagemodelname, torch_dtype=torch.bfloat16)
-    flux_pipe.enable_model_cpu_offload()
-    model_size = sum(p.numel() * p.element_size() for p in flux_pipe.transformer.parameters())
-    #tokenizer_size = 0  # FLUX doesn't use a separate tokenizer
-    loaded_models[imagemodelname] = flux_pipe
+#     flux_pipe = FluxPipeline.from_pretrained(imagemodelname, torch_dtype=torch.bfloat16)
+#     flux_pipe.enable_model_cpu_offload()
+#     model_size = sum(p.numel() * p.element_size() for p in flux_pipe.transformer.parameters())
+#     #tokenizer_size = 0  # FLUX doesn't use a separate tokenizer
+#     loaded_models[imagemodelname] = flux_pipe
 
-    # Calculate memory usage
-    final_memory = torch.cuda.memory_allocated()
-    memory_used = final_memory - initial_memory
+#     # Calculate memory usage
+#     final_memory = torch.cuda.memory_allocated()
+#     memory_used = final_memory - initial_memory
 
-    loaded_models[imagemodelname] = [str(time.time()), memory_used]
+#     loaded_models[imagemodelname] = [str(time.time()), memory_used]
 
-    lastimagemodelnameinloadfunction = (imagemodelname, model_size) #, tokenizer_size)
-    print(f"Model and tokenizer {imagemodelname} loaded successfully")
-    print(f"Model size: {get_size_str(model_size)}")
-    #print(f"Tokenizer size: {get_size_str(tokenizer_size)}")
-    print(f"GPU memory used: {get_size_str(memory_used)}")
+#     lastimagemodelnameinloadfunction = (imagemodelname, model_size) #, tokenizer_size)
+#     print(f"Model and tokenizer {imagemodelname} loaded successfully")
+#     print(f"Model size: {get_size_str(model_size)}")
+#     #print(f"Tokenizer size: {get_size_str(tokenizer_size)}")
+#     print(f"GPU memory used: {get_size_str(memory_used)}")
 
-    return (f"Model and tokenizer {imagemodelname} loaded successfully. "
-            f"Model size: {get_size_str(model_size)}, "
-            #f"Tokenizer size: {get_size_str(tokenizer_size)}, "
-            f"GPU memory used: {get_size_str(memory_used)}")
+#     return (f"Model and tokenizer {imagemodelname} loaded successfully. "
+#             f"Model size: {get_size_str(model_size)}, "
+#             #f"Tokenizer size: {get_size_str(tokenizer_size)}, "
+#             f"GPU memory used: {get_size_str(memory_used)}")
 
 
 def clear_all_models():
@@ -160,7 +160,7 @@ def loaded_model_list():
 
 # Initial model load
 load_model(modelname)
-load_image_model(imagemodelname)
+# load_image_model(imagemodelname)
 
 # Create embeddings for the knowledge base
 knowledge_base_embeddings = embedding_model.encode([doc["content"] for doc in knowledge_base])