Update app.py

app.py

@@ -1,13 +1,20 @@
 import gradio as gr
 import requests
 from PIL import Image
+import os 
 from transformers import BlipProcessor, BlipForConditionalGeneration
 import time
 from gradio_client import Client
 
+token = os.getenv('HF_TOKEN')
 blipper="Salesforce/blip-image-captioning-large"
 chatter="K00B404/transcript_image_generator"
 
+# Set your API endpoint and authorization details
+API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
+headers = {fAuthorization": f"Bearer {token}"}  # Replace with your actual token
+timeout = 60  # seconds
+
 # Load BLIP model for image captioning
 processor = BlipProcessor.from_pretrained(blipper)
 model = BlipForConditionalGeneration.from_pretrained(blipper)
@@ -32,9 +39,9 @@ Some characters might be terse while others might be more verbose."""
     return persona
 
 
-def dramaturg(message, max_tokens=256, temperature=0.5, top_p=0.95):
+def helper_llm(message, system_prompt, max_tokens=256, temperature=0.5, top_p=0.95):
     """Function to interact with the chatbot API using the generated persona"""
-    system_message="You are a Expert Dramaturg and your task is to use the input persona information and write a 'Role' description as compact instuctions for the actor"
+    
     try:
         # Call the API with the current message and system prompt (persona)
         response = chatbot_client.predict(
@@ -78,8 +85,11 @@ def generate_persona(img, min_len, max_len, persona_detail_level):
     # Calculate processing time
     end = time.time()
     total_time = f"Processing time: {end - start:.2f} seconds"
-    
-    return caption, dramaturg(persona), total_time
+   
+    # dramaturg to mae a solid role for a actor from pragmatic description
+    system_message="You are a Expert Dramaturg and your task is to use the input persona information and write a 'Role' description as compact instuctions for the actor"
+    persona = helper_llm(persona, system_prompt=system_prompt)
+    return caption, persona, total_time
 
 
 def chat_with_persona(message, history, system_message, max_tokens, temperature, top_p):
@@ -98,6 +108,48 @@ def chat_with_persona(message, history, system_message, max_tokens, temperature,
     except Exception as e:
         return f"Error communicating with the chatbot API: {str(e)}"
 
+
+
+
+
+def generate_flux_image(final_prompt, is_negative, steps, cfg_scale, seed, strength):
+    """
+    Generate an image using the FLUX model via Hugging Face's inference API.
+    The function sends a POST request with the given payload and returns the image,
+    along with the seed and prompt used.
+    """
+    payload = {
+        "inputs": final_prompt,
+        "is_negative": is_negative,
+        "steps": steps,
+        "cfg_scale": cfg_scale,
+        "seed": seed,
+        "strength": strength
+    }
+    
+    response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
+    
+    if response.status_code != 200:
+        print(f"Error: Failed to get image. Response status: {response.status_code}")
+        print(f"Response content: {response.text}")
+        if response.status_code == 503:
+            raise gr.Error(f"{response.status_code} : The model is being loaded")
+        raise gr.Error(f"{response.status_code}")
+    
+    try:
+        image_bytes = response.content
+        image = Image.open(io.BytesIO(image_bytes))
+        # Optionally save the image to a file (filename based on seed)
+        output_path = f"./output_{seed}.png"
+        image.save(output_path)
+        print(f'\033[1mGeneration completed!\033[0m (Prompt: {final_prompt})')
+        return output_path, str(seed), final_prompt
+    except Exception as e:
+        print(f"Error when trying to open the image: {e}")
+        return None, None, None
+
+
+        
 # Create Gradio interface with tabs
 with gr.Blocks(title="Image Character Persona Generator") as iface:
     # Store the generated persona in a state variable to share between tabs
@@ -129,6 +181,7 @@ with gr.Blocks(title="Image Character Persona Generator") as iface:
             3. Select detail level for the persona
             4. Click "Generate Character Persona"
             5. Switch to the "Test Persona" tab to chat with your character
+            6. create similar images inspired by the 'role' 
             """)
         
         # Second tab: Test Character Chat
@@ -176,6 +229,31 @@ with gr.Blocks(title="Image Character Persona Generator") as iface:
                       [msg, chatbot])
             
             clear_btn.click(clear_chat, outputs=chatbot)
+
+
+        # New Tab 3: Flux Image Generation
+        with gr.Tab("Flux Image Generation"):
+            gr.Markdown("### Flux Image Generation")
+            final_prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Enter your prompt for Flux...")
+            is_negative = gr.Checkbox(label="Use Negative Prompt", value=False)
+            steps = gr.Slider(minimum=10, maximum=100, step=1, value=50, label="Steps")
+            cfg_scale = gr.Slider(minimum=1, maximum=20, step=1, value=7, label="CFG Scale")
+            seed = gr.Number(value=42, label="Seed")
+            strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.8, label="Strength")
+            
+            generate_button = gr.Button("Generate Flux Image")
+            output_image = gr.Image(label="Generated Image")
+            output_seed = gr.Textbox(label="Seed Used")
+            output_prompt = gr.Textbox(label="Prompt Used")
+            
+            generate_button.click(
+                fn=generate_flux_image, 
+                inputs=[final_prompt, is_negative, steps, cfg_scale, seed, strength], 
+                outputs=[output_image, output_seed, output_prompt]
+            )
+
+
+
     
     # Function to update system prompt in Test tab when persona is generated
     def update_persona_state(caption, persona, time_output):