Update app.py

app.py

@@ -13,6 +13,11 @@ from threading import Thread
 import psutil
 import gc
 import torch
+import numpy as np
+from PIL import Image
+import stable_diffusion_cpp as sdcpp
+import base64
+import io
 
 load_dotenv()
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
@@ -217,9 +222,11 @@ global_data = {
     'end': {},
     'llama_perf_context_print': {},
     'llm_load_print_meta': {},
-    'model_type': {}
+    'model_type': {},
+    'image_model': {}
 }
 
+
 model_configs = [
     {
         "repo_id": "Hjgugugjhuhjggg/testing_semifinal-Q2_K-GGUF",
@@ -231,57 +238,63 @@ model_configs = [
         "filename": "Llama-3.2-3B-Instruct-uncensored-Q2_K.gguf",
         "name": "Llama-3.2-3B-Instruct"
     },
-    {
-        "repo_id": "Ffftdtd5dtft/Mistral-Nemo-Instruct-2407-Q2_K-GGUF",
-        "filename": "mistral-nemo-instruct-2407-q2_k.gguf",
-        "name": "Mistral-Nemo-Instruct"
-    },
-    {
-        "repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-70B-Instruct-Q2_K-GGUF",
-        "filename": "meta-llama-3.1-70b-instruct-q2_k.gguf",
-        "name": "Meta-Llama-3.1-70B-Instruct"
-    },
-    {
-        "repo_id": "Ffftdtd5dtft/Qwen2-Math-72B-Instruct-Q2_K-GGUF",
-        "filename": "qwen2-math-72b-instruct-q2_k.gguf",
-        "name": "Qwen2-Math-72B-Instruct"
-    },
-    {
-        "repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-70B-Q2_K-GGUF",
-        "filename": "meta-llama-3.1-70b-q2_k.gguf",
-        "name": "Meta-Llama-3.1-70B"
-    }
+     {
+        "repo_id": "city96/FLUX.1-schnell-gguf",
+        "filename": "flux1-schnell-Q2_K.gguf",
+        "name": "flux1-schnell"
+     },
+    
 ]
 
 class ModelManager:
     def __init__(self):
         self.models = {}
+        self.image_model = None
 
     def load_model(self, model_config):
-        if model_config['name'] not in self.models:
-            try:
-                self.models[model_config['name']] = Llama.from_pretrained(
-                    repo_id=model_config['repo_id'],
-                    filename=model_config['filename'],
-                    use_auth_token=HUGGINGFACE_TOKEN,
-                    n_threads=20,
-                    use_gpu=False
-                )
-            except Exception as e:
-                pass
+        if model_config['name'] not in self.models and model_config['name'] != "flux1-schnell":
+           try:
+               self.models[model_config['name']] = Llama.from_pretrained(
+                  repo_id=model_config['repo_id'],
+                  filename=model_config['filename'],
+                  use_auth_token=HUGGINGFACE_TOKEN,
+                  n_threads=20,
+                  use_gpu=False
+               )
+           except Exception as e:
+              pass
+
+    def load_image_model(self, model_config):
+       try:
+          self.image_model = sdcpp.StableDiffusionCpp(
+              repo_id=model_config['repo_id'],
+              filename=model_config['filename'],
+              use_auth_token=HUGGINGFACE_TOKEN,
+              n_threads=20,
+              use_gpu=False
+          )
+       except Exception as e:
+         print(f"Error loading image model: {e}")
 
     def load_all_models(self):
         with ThreadPoolExecutor() as executor:
             for config in model_configs:
-                executor.submit(self.load_model, config)
-        return self.models
+                if config['name'] == "flux1-schnell":
+                   executor.submit(self.load_image_model, config)
+                else:
+                    executor.submit(self.load_model, config)
+        return self.models, self.image_model
+
 
 model_manager = ModelManager()
-global_data['models'] = model_manager.load_all_models()
+global_data['models'], global_data['image_model'] = model_manager.load_all_models()
 
 class ChatRequest(BaseModel):
     message: str
 
+class ImageRequest(BaseModel):
+    prompt: str
+
 def normalize_input(input_text):
     return input_text.strip()
 
@@ -305,6 +318,7 @@ def cache_response(func):
         return response
     return wrapper
 
+
 @cache_response
 def generate_model_response(model, inputs):
     try:
@@ -332,9 +346,33 @@ async def process_message(message):
             for model_name, future in zip(global_data['models'].keys(), as_completed(futures))
         ]
     unique_responses = remove_repetitive_responses(responses)
-    formatted_response = next(iter(unique_responses.values()))  # Get a single response
+    formatted_response = next(iter(unique_responses.values()))
     return formatted_response
 
+async def generate_image(prompt: str):
+    if global_data['image_model']:
+        try:
+            image_bytes = global_data['image_model'].generate(
+                prompt=prompt,
+                negative_prompt="ugly, deformed, disfigured",
+                steps=25,
+                cfg_scale=7.0,
+                width=512,
+                height=512,
+                seed=-1,
+                return_type='bytes'
+             )
+             
+            image = Image.open(io.BytesIO(image_bytes))
+            return image
+        except Exception as e:
+           print(f"Error generating image: {e}")
+           return None
+    else:
+         print("No image model loaded.")
+         return None
+    
+
 app = FastAPI()
 
 @app.post("/generate")
@@ -345,6 +383,22 @@ async def generate(request: ChatRequest):
     except Exception as e:
         return JSONResponse(content={"error": str(e)})
 
+@app.post("/generate_image")
+async def generate_image_endpoint(request: ImageRequest):
+    try:
+        image = await generate_image(request.prompt)
+        if image:
+            buffered = io.BytesIO()
+            image.save(buffered, format="PNG")
+            image_base64 = base64.b64encode(buffered.getvalue()).decode()
+
+            return JSONResponse(content={"image": image_base64})
+        else:
+            return JSONResponse(content={"error": "Image generation failed or no model loaded"})
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)})
+
+
 def run_uvicorn():
     try:
         uvicorn.run(app, host="0.0.0.0", port=7860)
@@ -355,14 +409,28 @@ iface = gr.Interface(
     fn=process_message,
     inputs=gr.Textbox(lines=2, placeholder="Enter your message here..."),
     outputs=gr.Markdown(),
-    title="Multi-Model LLM API (CPU Optimized)",
+    title="Multi-Model LLM & Image API (CPU Optimized)",
     description="Optimized version using GPU and memory management techniques."
 )
+iface_image = gr.Interface(
+    fn=generate_image,
+    inputs=gr.Textbox(lines=2, placeholder="Enter image prompt here..."),
+    outputs=gr.Image(),
+    title="Stable Diffusion Image Generator",
+    description="Generate images using the specified stable diffusion model."
+)
+
 
 def run_gradio():
-    iface.launch(server_port=7862, prevent_thread_lock=True)
+    with gr.Blocks(title="Multi-Model LLM & Image API (CPU Optimized)") as demo:
+        with gr.Tab("LLM"):
+            iface.render()
+        with gr.Tab("Image Generator"):
+            iface_image.render()
+    demo.launch(server_port=7862, prevent_thread_lock=True)
+
 
 if __name__ == "__main__":
     Thread(target=run_uvicorn).start()
     Thread(target=run_gradio).start()
-    asyncio.get_event_loop().run_forever()
+    asyncio.get_event_loop().run_forever()