Upload 2 files

main.py

@@ -0,0 +1,291 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from langchain.memory import ConversationTokenBufferMemory
+import os
+import json
+import requests
+import time
+from PIL import Image
+from io import BytesIO
+from dotenv import load_dotenv
+import tempfile
+from fastapi.responses import FileResponse
+
+
+# Load environment variables
+load_dotenv()
+
+# Initialize FastAPI app
+app = FastAPI(title="Advanced AI Mock-up FastAPI")
+
+# Configure API Keys and global dependencies
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+url = os.getenv("IMAGE_API_URL")
+API_KEY = os.getenv("IMAGE_API_KEY")
+
+if not OPENAI_API_KEY or not API_KEY:
+    raise EnvironmentError("Missing API keys. Please set OPENAI_API_KEY and IMAGE_API_KEY in the environment variables.")
+
+from openai import OpenAI
+# Configure OpenAI client
+client = OpenAI()
+
+from langchain_openai import OpenAI
+llm = OpenAI()
+memory = ConversationTokenBufferMemory(llm=llm, max_token_limit=4000)
+
+# API Key and Headers for image generation
+headers = {
+    "accept": "application/json",
+    "x-key": API_KEY,
+    "Content-Type": "application/json"
+}
+
+# Pydantic model for input
+class ConversationRequest(BaseModel):
+    question: str
+
+# Function to manage greeting
+def Greeting(question, chat_history):
+    prompt = f"""
+    You are a professional AI assistant specialized in AI-powered mock-up creation. Start with a warm greeting, ask about the user's well-being, and also ask related to AI-powered mock-up creation for jackets or other apparel. Tailor your conversation to establish a friendly and professional tone.
+
+    Chat History:
+    {chat_history}
+    """
+
+    response = client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": prompt},
+            {"role": "user", "content": f"Question: {question}"}
+        ]
+    )
+    return response.choices[0].message.content
+
+def select_state(chat_history):
+    output_format = '''
+    Answer according to the following JSON format:
+    {
+    "State": "Here you will select one state based on chat history: 'greeting', 'gather_info', 'analyze_chat_history', 'generate_images'"
+    }'''
+
+    prompt = f"""
+    Based on the below chat history, decide the state for the agent. The state can be:
+    - 'greeting': if the chat history lacks a greeting message.
+    - 'gather_info': if greeting messages (like 'hi', 'hello', 'how are you') have been successfully executed.
+    - 'analyze_chat_history': if sufficient information has been gathered, including:
+        - Team name.
+        - Colors or style preferences.
+        - Details about patterns, or any unique requirements.
+    - 'generate_images': if image prompts are generated.
+
+    Chat History:
+    {chat_history}
+    """ + output_format
+
+    response = client.chat.completions.create(
+        model="gpt-4o",
+        response_format={"type": "json_object"},
+        messages=[
+            {"role": "system", "content": prompt},
+            {"role": "user", "content": "Select the next state"}
+        ]
+    )
+    json_data = json.loads(response.choices[0].message.content)
+    return json_data['State']
+
+
+# Function to gather information
+def Gather_info(question, chat_history):
+    prompt = f"""
+    You are an information-gathering agent specialized in AI-powered mock-up creation. Your task is to politely gather the following information from the user:
+
+    - Team: Ask what team this is for.
+    - Team colors: Ask for the team colors or other specific colors they want to use.
+    - Style guide: Inquire if the user can provide a details to the team style guide.
+    - Additional details: Gather any additional specific information related to the Team, such as patterns, or any unique requirements.
+
+    Please ask these questions one by one in a friendly and engaging manner, and ensure you document all the provided details accurately.
+
+    Chat History:
+    {chat_history}
+    """
+
+    response = client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": prompt},
+            {"role": "user", "content": f"Question: {question}"}
+        ]
+    )
+    return response.choices[0].message.content
+
+def analyze_chat_history(chat_history):
+    output_format = '''
+    Answer according to the following JSON format:
+    {
+    "Analysis": "Provide a summary analysis of the chat history, focusing on key insights derived from the gathered information.",
+    "NextAction": "Specify the next logical action: either continue the conversation or conclude it.",
+    "Prompts": [
+        "Prompt 1: Detailed prompt for generating the first mock-up.",
+        "Prompt 2: Detailed prompt for generating the second mock-up.",
+        "Prompt 3: Detailed prompt for generating the third mock-up.",
+        "Prompt 4: Detailed prompt for generating the fourth mock-up."
+    ]
+    }'''
+
+    prompt = f"""
+You are a highly intelligent and efficient analysis agent tasked with processing the chat history provided below. Based solely on the relevant information gathered by the information-gathering agent, your responsibilities are to:
+
+1. Summarize the user's key points and design requirements with precision, highlighting the essential elements.
+2. Generate 4 detailed and creative prompts for image mock-ups tailored to the user's specific needs.
+3. In all the prompts the information about the jacket should be same so jacket in all the images are same but have different view.
+
+Ensure that the generated prompts adhere to the following criteria:
+- Visually compelling, emphasizing creativity, detail, and storytelling.
+- Highly specific, incorporating the following aspects where applicable:
+  - Key themes, team dynamics, or user-specified concepts.
+  - Color schemes, textures, and style guidelines.
+  - Camera and Lens Settings: Recommend camera models (e.g., Canon EOS R5, Nikon Z9), lenses (e.g., 50mm f/1.8 for portraits or 85mm for close-ups), and techniques (e.g., shallow depth of field, macro for texture).
+  - Artistic Enhancements: Suggest details like angles (e.g., low-angle, top-down), effects (e.g., bokeh, soft focus), or scene accents (e.g., props or natural textures).
+  - Aspect Ratio and Style Tags: Specify dimensions (e.g., --ar 16:9 for banners or --ar 4:5 for Instagram). Include style tags like --style cinematic, --style raw, or --style editorial.
+  - Lighting details, including time of day, intensity, direction, and color temperature.
+  - Composition elements like framing, depth of field, symmetry, and rule of thirds.
+  - Environmental and contextual details that provide additional realism or artistic flair.
+- Clearly structured to provide effective guidance for advanced image generation models.
+- Prompts should Focus on the provided color combination. Do not add anything from yourself use all the context that user have provided
+- Do not add Humans in the images. Only generate the images of the jackets in the white background
+Chat History:
+{chat_history}
+    """ + output_format
+
+    response = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {"role": "system", "content": prompt},
+            {"role": "user", "content": "Analyze the conversation and generate prompts"}
+        ]
+    )
+
+    # Extract response content
+    response_content = response.choices[0].message.content.strip()
+
+    # Clean and validate response content
+    if response_content.startswith("```") and response_content.endswith("```"):
+        response_content = response_content[response_content.find("\n") + 1 : -3].strip()
+
+    if not response_content:
+        raise ValueError("The API returned an empty response.")
+
+    try:
+        json_data = json.loads(response_content)
+    except json.JSONDecodeError as e:
+        print("Error parsing JSON:", e)
+        print("Content causing error:", response_content)
+        raise
+
+    return json_data
+
+# Create a temporary directory to store generated images
+temp_dir = tempfile.TemporaryDirectory()
+
+def generate_images(prompts, url, headers):
+    temp_dir = tempfile.TemporaryDirectory()
+    image_links = []
+    
+    for index, prompt in enumerate(prompts, start=1):
+        print(f"Generating image {index} of {len(prompts)}")
+
+        payload = {
+            "prompt": prompt,
+            "width": 1024,
+            "height": 1024,
+            "guidance_scale": 1,
+            "num_inference_steps": 50,
+            "max_sequence_length": 512,
+        }
+
+        response = requests.post(url, headers=headers, json=payload).json()
+        if "id" not in response:
+            print("Error in generating image:", response)
+            continue
+
+        request_id = response["id"]
+        print(f"Image generation request ID for prompt {index}: {request_id}")
+
+        while True:
+            time.sleep(0.5)
+            result = requests.get(
+                "https://api.bfl.ml/v1/get_result",
+                headers=headers,
+                params={"id": request_id},
+            ).json()
+
+            if result["status"] == "Ready":
+                if "result" in result and "sample" in result["result"]:
+                    image_url = result["result"]["sample"]
+                    print(f"Generated image URL for prompt {index}: {image_url}")
+                    image_links.append(image_url)
+                else:
+                    print(f"Error: 'sample' key not found in the result for prompt {index}.")
+                break
+            else:
+                print(f"Image generation status for prompt {index}: {result['status']}")
+    
+    return image_links
+
+
+def manage_conversation(question, url, headers, memory):
+    chat_history = memory.load_memory_variables({})
+    chat_history = chat_history['history']
+
+    # Get the current state
+    state = select_state(chat_history)
+
+    if state == "greeting":
+        response = Greeting(question, chat_history)
+    elif state == "gather_info":
+        response = Gather_info(question, chat_history)
+    elif state == "analyze_chat_history":
+        response = analyze_chat_history(chat_history)
+        # Serialize the JSON response to a string if it's a dictionary
+        response = json.dumps(response, indent=4)
+    elif state == "generate_images":
+        prompts = analyze_chat_history(chat_history)['Prompts']
+        image_links = generate_images(prompts, url=url, headers=headers)
+        response = json.dumps({"message": "Images generated successfully.", "image_links": image_links}, indent=4)
+    else:
+        response = "Conversation ended."
+
+    # Save the response to memory as a string
+    memory.save_context({"input": question}, {"output": response})
+
+    return response
+
+
+# API Endpoint
+@app.post("/conversation/")
+async def conversation_endpoint(request: ConversationRequest):
+    try:
+        response = manage_conversation(request.question, url, headers, memory)  # Pass the required arguments
+        return {"response": response}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/new_chat/")
+async def new_chat():
+    """
+    This endpoint resets the memory and starts a new chat session.
+    """
+    try:
+        global memory
+        memory = ConversationTokenBufferMemory(llm=llm, max_token_limit=4000)
+        return {"message": "New chat session started successfully."}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/")
+async def root():
+    return {"message": "API is up and running!"}

requirements.txt

@@ -0,0 +1,3 @@
+openai
+langchain
+langchain_openai