Update app.py

app.py

@@ -1,171 +1,169 @@
-import os
-import threading
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
-from fastapi.responses import  FileResponse
-from fastapi.staticfiles import StaticFiles
-from huggingface_hub import InferenceClient
-from concurrent.futures import ThreadPoolExecutor
-from fastapi.responses import HTMLResponse
-from groq import Groq
-import uvicorn
-import base64
-from io import BytesIO
-import os
-import requests
-
-API_URL = "https://api-inference.huggingface.co/models/Falconsai/nsfw_image_detection"
-headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
-
-def is_image_safe(data):
-    try:
-        response = requests.post(API_URL, headers=headers, data=data)
-        response.raise_for_status()  # Raise an error for HTTP issues
-        result = response.json()
-        print(result)
-
-        # Look for the 'nsfw' label and extract its score
-        nsfw_score = next((item["score"] for item in result if item["label"] == "nsfw"), None)
-
-        # Ensure the 'nsfw' label exists in the response
-        if nsfw_score is not None:
-            return nsfw_score < 0.5  # Return True if the NSFW score is less than 0.5
-        else:
-            print("Error: 'nsfw' label not found in response.")
-            return False
-    except requests.exceptions.RequestException as e:
-        print(f"Error: Failed to process the request. {e}")
-        return False
-
-
-app = FastAPI()
-
-# Serve static files (HTML, CSS, JS)
-app.mount("/static", StaticFiles(directory="static"), name="static")
-
-# Initialize a single Groq client and inference client
-client = Groq()
-image_gen_client = InferenceClient("black-forest-labs/FLUX.1-schnell")
-executor = ThreadPoolExecutor()
-
-@app.get("/", response_class=FileResponse)
-async def get():
-    return FileResponse("static/index.html")
-
-
-@app.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
-
-    def generate_text(input_text):
-        chat_completion = client.chat.completions.create(
-            messages=[
-                {"role": "system", "content": "You are an assistant that generates gratitude journal entries. Focus solely on expressing gratitude in a concise, meaningful tone. Do not include introductory or concluding statements. Avoid unrelated topics or personal opinions. Keep the entries simple, neutral, and directly reflective of gratitude."},
-                {"role": "user", "content": input_text},
-            ],
-            model="llama3-8b-8192",
-        )
-        print(chat_completion.choices[0].message.content)
-        return chat_completion.choices[0].message.content
-    
-    def generate_image_prompt(input_text):
-        chat_completion = client.chat.completions.create(
-            messages=[
-                {"role": "system", "content": "You are an assistant that generates image prompts reflecting gratitude journal entries. Create concise, vivid descriptions that visually represent gratitude and positive moments. Avoid adding any introductory or concluding statements. Keep the prompts simple, neutral, and focused solely on the imagery."},
-                {"role": "user", "content": input_text},
-            ],
-            model="llama3-8b-8192",
-        )
-        print(chat_completion.choices[0].message.content)
-        return chat_completion.choices[0].message.content
-
-    def analyze_image(image_data):
-        completion = client.chat.completions.create(
-            model="llama-3.2-11b-vision-preview",
-            messages=[
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "text", "text": "Describe the contents of this image for a gratitude journal."},
-                        {"type": "image", "image": image_data},
-                    ],
-                }
-            ],
-            temperature=1,
-            max_tokens=1024,
-            top_p=1,
-            stream=False,
-            stop=None,
-        )
-        print(completion.choices[0].message.content)
-        return completion.choices[0].message.content
-
-    def handle_text_to_speech(audio_data):
-        transcription = client.audio.transcriptions.create(
-            file=("audio.mp3", audio_data),
-            model="whisper-large-v3-turbo",
-            response_format="json",
-            language="en",
-        )
-        return transcription.text
-
-    def generate_image(prompt):
-        print(prompt)
-        return image_gen_client.text_to_image(prompt)
-
-    try:
-        while True:
-            data = await websocket.receive_json()
-            input_text = data.get("text", "")
-            input_images = data.get("images", [])
-            input_audio = data.get("audio", None)
-
-            journal_text = input_text
-            tasks = []
-
-            # Process images
-            for image in input_images:
-                def process_image(image=image):
-                    image_description = analyze_image(image)
-                    nonlocal journal_text
-                    journal_text += f" {image_description}"  
-                tasks.append(executor.submit(process_image))
-
-            # Process audio
-            if input_audio:
-                def process_audio():
-                    audio_transcription = handle_text_to_speech(input_audio)
-                    nonlocal journal_text
-                    journal_text += f" {audio_transcription}"  
-                tasks.append(executor.submit(process_audio))
-
-            # Wait for all tasks to complete
-            for task in tasks:
-                task.result()
-
-            # Generate journal entry
-            journal_output = generate_text(journal_text)
-            journal_image_prompt= generate_image_prompt(journal_text)
-            await websocket.send_json({"type": "journal", "content": journal_output})
-            await websocket.send_json({"type": "info", "content": "Gratitude-themed image is being generated"})
-            image_prompt = f"Generate a gratitude-themed artistic image based on this journal entry: {journal_image_prompt}"
-            generated_image = generate_image(image_prompt)  # This returns a PIL.Image object
-            print("image generated")
-            
-
-
-            # Convert PIL.Image to Base64
-            buffered = BytesIO()
-            generated_image.save(buffered, format="PNG")  # Save as PNG or any format you prefer
-            img_str=""
-            if is_image_safe(buffered.getvalue()):
-                img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")  # Base64 encode the image
-
-            # Send the image as a base64 string via WebSocket
-            await websocket.send_json({"type": "image", "image": img_str})
-
-            
-
-    except WebSocketDisconnect:
-        await websocket.close()
-if __name__ == "__main__":
-    uvicorn.run("app:app", host="0.0.0.0", port=8090, reload=True)
+import os
+import threading
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.responses import  FileResponse
+from fastapi.staticfiles import StaticFiles
+from huggingface_hub import InferenceClient
+from concurrent.futures import ThreadPoolExecutor
+from fastapi.responses import HTMLResponse
+from groq import Groq
+import uvicorn
+import base64
+from io import BytesIO
+import os
+import requests
+
+API_URL = "https://api-inference.huggingface.co/models/Falconsai/nsfw_image_detection"
+headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
+
+def is_image_safe(data):
+    try:
+        response = requests.post(API_URL, headers=headers, data=data)
+        response.raise_for_status()  # Raise an error for HTTP issues
+        result = response.json()
+        print(result)
+
+        # Look for the 'nsfw' label and extract its score
+        nsfw_score = next((item["score"] for item in result if item["label"] == "nsfw"), None)
+
+        # Ensure the 'nsfw' label exists in the response
+        if nsfw_score is not None:
+            return nsfw_score < 0.5  # Return True if the NSFW score is less than 0.5
+        else:
+            print("Error: 'nsfw' label not found in response.")
+            return False
+    except requests.exceptions.RequestException as e:
+        print(f"Error: Failed to process the request. {e}")
+        return False
+
+
+app = FastAPI()
+
+# Serve static files (HTML, CSS, JS)
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+# Initialize a single Groq client and inference client
+client = Groq()
+image_gen_client = InferenceClient("black-forest-labs/FLUX.1-schnell")
+executor = ThreadPoolExecutor()
+
+@app.get("/", response_class=FileResponse)
+async def get():
+    return FileResponse("static/index.html")
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+
+    def generate_text(input_text):
+        chat_completion = client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": "You are an assistant that generates gratitude journal entries. Focus solely on expressing gratitude in a concise, meaningful tone. Do not include introductory or concluding statements. Avoid unrelated topics or personal opinions. Keep the entries simple, neutral, and directly reflective of gratitude."},
+                {"role": "user", "content": input_text},
+            ],
+            model="llama3-8b-8192",
+        )
+        print(chat_completion.choices[0].message.content)
+        return chat_completion.choices[0].message.content
+    
+    def generate_image_prompt(input_text):
+        chat_completion = client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": "You are an assistant that generates image prompts reflecting gratitude journal entries. Create concise, vivid descriptions that visually represent gratitude and positive moments. Avoid adding any introductory or concluding statements. Keep the prompts simple, neutral, and focused solely on the imagery."},
+                {"role": "user", "content": input_text},
+            ],
+            model="llama3-8b-8192",
+        )
+        print(chat_completion.choices[0].message.content)
+        return chat_completion.choices[0].message.content
+
+    def analyze_image(image_data):
+        completion = client.chat.completions.create(
+            model="llama-3.2-11b-vision-preview",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "Describe the contents of this image for a gratitude journal."},
+                        {"type": "image", "image": image_data},
+                    ],
+                }
+            ],
+            temperature=1,
+            max_tokens=1024,
+            top_p=1,
+            stream=False,
+            stop=None,
+        )
+        print(completion.choices[0].message.content)
+        return completion.choices[0].message.content
+
+    def handle_text_to_speech(audio_data):
+        transcription = client.audio.transcriptions.create(
+            file=("audio.mp3", audio_data),
+            model="whisper-large-v3-turbo",
+            response_format="json",
+            language="en",
+        )
+        return transcription.text
+
+    def generate_image(prompt):
+        print(prompt)
+        return image_gen_client.text_to_image(prompt)
+
+    try:
+        while True:
+            data = await websocket.receive_json()
+            input_text = data.get("text", "")
+            input_images = data.get("images", [])
+            input_audio = data.get("audio", None)
+
+            journal_text = input_text
+            tasks = []
+
+            # Process images
+            for image in input_images:
+                def process_image(image=image):
+                    image_description = analyze_image(image)
+                    nonlocal journal_text
+                    journal_text += f" {image_description}"  
+                tasks.append(executor.submit(process_image))
+
+            # Process audio
+            if input_audio:
+                def process_audio():
+                    audio_transcription = handle_text_to_speech(input_audio)
+                    nonlocal journal_text
+                    journal_text += f" {audio_transcription}"  
+                tasks.append(executor.submit(process_audio))
+
+            # Wait for all tasks to complete
+            for task in tasks:
+                task.result()
+
+            # Generate journal entry
+            journal_output = generate_text(journal_text)
+            journal_image_prompt= generate_image_prompt(journal_text)
+            await websocket.send_json({"type": "journal", "content": journal_output})
+            await websocket.send_json({"type": "info", "content": "Gratitude-themed image is being generated"})
+            image_prompt = f"Generate a gratitude-themed artistic image based on this journal entry: {journal_image_prompt}"
+            generated_image = generate_image(image_prompt)  # This returns a PIL.Image object
+            print("image generated")
+            
+
+
+            # Convert PIL.Image to Base64
+            buffered = BytesIO()
+            generated_image.save(buffered, format="PNG")  # Save as PNG or any format you prefer
+            img_str=""
+            if is_image_safe(buffered.getvalue()):
+                img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")  # Base64 encode the image
+
+            # Send the image as a base64 string via WebSocket
+            await websocket.send_json({"type": "image", "image": img_str})
+
+            
+
+    except WebSocketDisconnect:
+        await websocket.close()