Update app.py

app.py

@@ -15,7 +15,7 @@ import pytesseract
 from fontTools.ttLib import TTFont
 from langchain_groq import ChatGroq
 import logging
-from huggingface_hub import InferenceClient
+from together import Together
 
 
 # Load environment variables
@@ -100,7 +100,7 @@ def advertisement_generator():
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-client = InferenceClient("black-forest-labs/FLUX.1-dev", token=os.getenv('HF_TOKEN'))
+client = Together()
 
 def generate_poster():
     #st.header("Generate Social Media Post")
@@ -128,11 +128,11 @@ def generate_poster():
         if aspect_ratio == "1:1":
             dimensions = st.selectbox("Select Image Dimensions", ["1024x1024", "1200x1200", "1504x1504"])
         elif aspect_ratio == "16:9":
-            dimensions = st.selectbox("Select Image Dimensions", ["1920x1080", "1280x720", "2560x1440"])
+            dimensions = st.selectbox("Select Image Dimensions", ["1024x576", "1280x720", "1792x1008"])
         elif aspect_ratio == "4:5":
-            dimensions = st.selectbox("Select Image Dimensions", ["1080x1344", "1200x1504", "1600x2000"])
+            dimensions = st.selectbox("Select Image Dimensions", ["1024x1280", "1200x1500", "1600x2000"])
         elif aspect_ratio == "9:16":
-            dimensions = st.selectbox("Select Image Dimensions", ["1080x1920", "1344x2400", "1600x2848"])
+            dimensions = st.selectbox("Select Image Dimensions", ["576x1024", "720x1280", "1008x1792"])
         
         design_style = st.selectbox("Select Design Style", [
             "Minimalistic", "Bold/Graphic", "Elegant", "Playful/Fun",
@@ -174,8 +174,11 @@ def generate_poster():
 
     if generate_button:
         # Generate 4 different variations of the prompt with enhanced realism and attention-grabbing elements
+        lighting_options = ['golden hour lighting', 'studio lighting', 'natural daylight', 'dramatic spotlights']
+        visual_elements = ['3D elements', 'metallic accents', 'glass effects', 'neon highlights']
+        
         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)}"
+            f"Create a professional and eye-catching {post_type.lower()} advertisement. The image should feature impactful {selected_color.lower() if selected_color != 'None' else 'vibrant'} colors that align with brand identity. Header: \"{header}\". Sub-header: \"{sub_header}\". Implement a {design_style.lower()} design style with clean, commercial-grade visuals. Main focus: {description}. Compose in {aspect_ratio} aspect ratio at {width}x{height}. The design should incorporate modern {random.choice(visual_elements)} to enhance visual appeal. Ensure high resolution with perfect clarity and legibility. Text should be bold, clear and strategically placed for maximum impact. Create a compelling visual hierarchy that drives attention to key messaging and call-to-action elements. Make it look like a premium advertisement created by a professional design agency. Variation {i+1}/4."
             for i in range(4)
         ]
 
@@ -184,11 +187,36 @@ def generate_poster():
             with st.spinner(f"Generating Graphic {i+1}..."):
                 logger.info(f"Generating Graphic {i+1} with prompt: {prompt}")
                 
-                # Generate image using FLUX model with random seed
-                random_seed = random.randint(1, 999999)
-                image = client.text_to_image(prompt, seed=random_seed)
+                # Adjust dimensions if needed to stay within API limits
+                adjusted_width = min(1792, max(64, width))
+                adjusted_height = min(1792, max(64, height))
+                
+                # Maintain aspect ratio while adjusting dimensions
+                if width > 1792 or height > 1792:
+                    ratio = min(1792/width, 1792/height)
+                    adjusted_width = int(width * ratio)
+                    adjusted_height = int(height * ratio)
                 
-                if image:
+                # Generate image using Together API
+                response = client.images.generate(
+                    prompt=prompt,
+                    model="black-forest-labs/FLUX.1-schnell-Free",
+                    width=adjusted_width,
+                    height=adjusted_height,
+                    steps=4,
+                    n=1,
+                    response_format="b64_json"
+                )
+                
+                if response.data:
+                    # Convert base64 to image
+                    image_data = base64.b64decode(response.data[0].b64_json)
+                    image = Image.open(io.BytesIO(image_data))
+                    
+                    # Resize back to original dimensions if needed
+                    if adjusted_width != width or adjusted_height != height:
+                        image = image.resize((width, height), Image.LANCZOS)
+                    
                     # Add logo if provided
                     if logo:
                         logo_image = Image.open(logo)
@@ -274,9 +302,21 @@ def generate_image_prompt(image):
     return response.text
 
 def generate_new_image(prompt):
-    # Generate image using FLUX model with random seed
-    random_seed = random.randint(1, 999999)
-    return client.text_to_image(prompt, seed=random_seed)
+    # Generate image using Together API
+    response = client.images.generate(
+        prompt=prompt,
+        model="black-forest-labs/FLUX.1-schnell-Free",
+        width=1024,
+        height=768,
+        steps=4,
+        n=1,
+        response_format="b64_json"
+    )
+    
+    if response.data:
+        image_data = base64.b64decode(response.data[0].b64_json)
+        return Image.open(io.BytesIO(image_data))
+    return None
 
 # Part 2: Image Editing