Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from PIL import Image
+from PIL import Image, ImageFont, ImageDraw
 import io
 import base64
 import google.generativeai as genai
@@ -15,6 +15,7 @@ import pytesseract
 from fontTools.ttLib import TTFont
 from langchain_groq import ChatGroq
 import logging
+from huggingface_hub import InferenceClient
 
 
 # Load environment variables
@@ -36,11 +37,8 @@ model = genai.GenerativeModel(
     generation_config=generation_config,
 )
 
-API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
 
-# Set up Tesseract path
-pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'  # Adjust this path as needed
+pytesseract.pytesseract.tesseract_cmd = r'Tesseract-OCR\tesseract.exe'
 
 # Set up Groq LLM
 llm = ChatGroq(
@@ -49,6 +47,7 @@ llm = ChatGroq(
     model_name="llama-3.1-70b-versatile"
 )
 
+
 # Content from agent.py
 def generate_advertisement_prompt(description):
     prompt = f"""
@@ -87,38 +86,21 @@ def advertisement_generator():
 
     post_description = st.text_input("Enter a brief description for your advertisement post:")
 
-    if st.button("Generate Advertisement Elements"):
+    if st.button("Generate Image Prompt"):
         if post_description:
-            with st.spinner("Generating advertisement elements..."):
+            with st.spinner("Prompt Enhancer..."):
                 final_prompt = generate_advertisement_prompt(post_description)
-                st.subheader("Generated Advertisement Prompt:")
+                st.subheader("Generated Image Prompt:")
                 
                 st.text_area(label="Final Prompt", value=final_prompt.strip(), height=200)
         else:
-            st.warning("Please enter a description for your advertisement post.")
+            st.warning("Please enter a description for your post.")
 
 # Content from image_generation.py
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-def query(payload):
-    try:
-        response = requests.post(API_URL, headers=headers, json=payload)
-        response.raise_for_status()
-        
-        logger.debug(f"API response status code: {response.status_code}")
-        logger.debug(f"API response headers: {response.headers}")
-        
-        return response.content
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Request failed: {str(e)}")
-        st.error(f"Request failed: {str(e)}")
-        return None
-
-def increase_image_quality(image, scale_factor):
-    width, height = image.size
-    new_size = (width * scale_factor, height * scale_factor)
-    return image.resize(new_size, Image.LANCZOS)
+client = InferenceClient("black-forest-labs/FLUX.1-dev", token=os.getenv('HF_TOKEN'))
 
 def generate_poster():
     #st.header("Generate Social Media Post")
@@ -133,7 +115,7 @@ def generate_poster():
     st.markdown("<div style='height: 30px;'></div>", unsafe_allow_html=True)
 
     if st.session_state.get('show_ad_generator', False):
-        with st.expander("Advertisement Generator :arrow_right:", expanded=True):
+        with st.expander("Prompt Enhancer :arrow_right:", expanded=True):
             advertisement_generator()
 
     st.markdown("<div style='height: 30px;'></div>", unsafe_allow_html=True)
@@ -144,13 +126,13 @@ def generate_poster():
         aspect_ratio = st.selectbox("Select Image Aspect Ratio", ["1:1", "16:9", "4:5", "9:16"])
     with col2:
         if aspect_ratio == "1:1":
-            dimensions = st.selectbox("Select Image Dimensions", ["1080x1080", "1200x1200", "1500x1500"])
+            dimensions = st.selectbox("Select Image Dimensions", ["1024x1024", "1200x1200", "1504x1504"])
         elif aspect_ratio == "16:9":
             dimensions = st.selectbox("Select Image Dimensions", ["1920x1080", "1280x720", "2560x1440"])
         elif aspect_ratio == "4:5":
-            dimensions = st.selectbox("Select Image Dimensions", ["1080x1350", "1200x1500", "1600x2000"])
+            dimensions = st.selectbox("Select Image Dimensions", ["1080x1344", "1200x1504", "1600x2000"])
         elif aspect_ratio == "9:16":
-            dimensions = st.selectbox("Select Image Dimensions", ["1080x1920", "1350x2400", "1600x2844"])
+            dimensions = st.selectbox("Select Image Dimensions", ["1080x1920", "1344x2400", "1600x2848"])
         
         design_style = st.selectbox("Select Design Style", [
             "Minimalistic", "Bold/Graphic", "Elegant", "Playful/Fun",
@@ -190,35 +172,23 @@ def generate_poster():
 
     st.markdown("<div style='height: 30px;'></div>", unsafe_allow_html=True)
 
-    # Add number of images to generate (hidden from user)
-    #num_images = st.number_input("Number of images to generate", min_value=4, max_value=100, value=1)
-
-    num_images = 4
-
     if generate_button:
-        base_prompt = f"Create a minimalistic and realistic poster image with the following text as the header: '{header}'. Use '{sub_header}' as the sub-header. Use this design style and go with this {design_style.lower()} style for the post. Include the following text as the description for the advertisement: {description}. The image should have an aspect ratio of {aspect_ratio} and dimensions of {width}x{height}. Use {selected_color.lower()} as the dominant color in the image. Ensure the image is in 4k high resolution. And do not include any extra text in the image only instruction text included.And try to fill the vacant part with releveng images. Make the design clean and elegant, focusing on simplicity and realism."
-
-        if post_type == "Other":
-            full_prompt = f"{base_prompt} The image should be suitable for general use across various platforms."
-        else:
-            full_prompt = f"{base_prompt} This image is specifically for a {post_type.lower()}."
-
-        # Add user prompts to full_prompt if provided
-        for i, user_prompt in enumerate(user_prompts):
-            if user_prompt:
-                full_prompt += f" Include the following text in the image (Prompt {i+1}): '{user_prompt}'."
+        # Generate 4 different variations of the prompt with enhanced realism and attention-grabbing elements
+        prompt_variations = [
+            f"Create a hyper-realistic, professional {post_type.lower()} that demands attention. The image should feature stunning {selected_color.lower() if selected_color != 'None' else 'vibrant'} colors with perfect lighting and shadows. Header: \"{header}\". Sub-header: \"{sub_header}\". Implement a {design_style.lower()} aesthetic with photorealistic textures and materials. Main focus: {description}. Use cinema-quality composition with {aspect_ratio} aspect ratio at {width}x{height}. Include dramatic lighting effects, subtle lens flares, and professional-grade depth of field. The design should incorporate cutting-edge {random.choice(['3D elements', 'metallic accents', 'glass effects', 'neon highlights'])} to create visual interest. Ensure 8K ultra-high resolution with perfect clarity. Text should be crystal clear and expertly integrated. Create a powerful visual hierarchy that guides the viewer's eye. Make it look like it was shot by a professional photographer with high-end equipment. Variation {i+1}/4 with unique {random.choice(['golden hour lighting', 'studio lighting', 'natural daylight', 'dramatic spotlights'])}. Seed: {random.randint(1, 999999)}"
+            for i in range(4)
+        ]
 
         generated_images = []
-        for i in range(num_images):
+        for i, prompt in enumerate(prompt_variations):
             with st.spinner(f"Generating Graphic {i+1}..."):
-                logger.info(f"Generating Graphic {i+1} with prompt: {full_prompt}")
-                # Add a random seed to ensure different images
-                seed = random.randint(1, 1000000)
-                image_bytes = query({"inputs": full_prompt, "parameters": {"seed": seed, "width": width, "height": height}})
+                logger.info(f"Generating Graphic {i+1} with prompt: {prompt}")
                 
-                if image_bytes:
-                    image = Image.open(io.BytesIO(image_bytes))
-                    
+                # Generate image using FLUX model with random seed
+                random_seed = random.randint(1, 999999)
+                image = client.text_to_image(prompt, seed=random_seed)
+                
+                if image:
                     # Add logo if provided
                     if logo:
                         logo_image = Image.open(logo)
@@ -304,14 +274,9 @@ def generate_image_prompt(image):
     return response.text
 
 def generate_new_image(prompt):
-    # Add a random seed to ensure different images
-    seed = random.randint(1, 1000000)
-    payload = {
-        "inputs": prompt,
-        "parameters": {"seed": seed}
-    }
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return Image.open(io.BytesIO(response.content))
+    # Generate image using FLUX model with random seed
+    random_seed = random.randint(1, 999999)
+    return client.text_to_image(prompt, seed=random_seed)
 
 # Part 2: Image Editing