app.py

from PIL import Image
import gradio as gr
import torch
from diffusers import FluxPipeline
from groq import Groq  # Import the Groq library
from cryptography.fernet import Fernet
from huggingface_hub import login
def get_hf_token(encrypted_token):
    # Retrieve the decryption key from an environment variable
    key = "K4FlQbffvTcDxT2FIhrOPV1eue6ia45FFR3kqp2hHbM="
    if not key:
        raise ValueError("Missing decryption key! Set the DECRYPTION_KEY environment variable.")
    
    # Convert key from string to bytes if necessary
    if isinstance(key, str):
        key = key.encode()
    
    f = Fernet(key)
    # Decrypt and decode the token
    decrypted_token = f.decrypt(encrypted_token).decode()
    return decrypted_token
groq_client = Groq(api_key="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")
decrypted_token = get_hf_token("gAAAAABn3GfShExoJd50nau3B5ZJNiQ9dRD1ACO3XXMwVaIQMkmi59cL-MKGr6SYnsB0E2gGITJG2j29Ar9yjaZP-EC6hHsCBmwKSj4aFtTor9_n0_NdMBv1GtlxZRmwnQwriB-Xr94e")
login(token=decrypted_token)
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16
)
pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power

def enhance_prompt(user_prompt):
    """Enhances the given prompt using Groq and returns the refined prompt."""
    try:
        chat_completion = groq_client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": (
                        "Create prompts that paint a clear picture for image generation. "
                        "Be precise, detailed, and direct, describing not only the content of the image "
                        "but also details like tone, style, color palette, and point of view. "
                        "For photorealistic images, include the device used (e.g., 'shot on iPhone 16'), "
                        "aperture, lens, and shot type. Use clear, visual descriptions. "
                        "Prompt Structure: 'A [medium] of [subject], [subject’s characteristics], "
                        "[relation to background] [background]. [Details of background] "
                        "[Interactions with color and lighting]. (\"Taken on:\"/\"Drawn with:\")[Specific traits of style]' "
                        "Include details on medium, subject (clothing, hairstyle, pose, etc.), background, "
                        "colors, lighting, style traits, influences, technique, and camera settings if applicable."
                    ),
                },
                {"role": "user", "content": user_prompt}
            ],
            model="llama-3.3-70b-versatile",
            temperature=0.5,
            max_completion_tokens=1024,
            top_p=1,
            stop=None,
            stream=False,
        )
        enhanced = chat_completion.choices[0].message.content
    except Exception as e:
        enhanced = f"Error enhancing prompt: {str(e)}"
    return enhanced

def generate_image(prompt):
    """Generates an image using the refined prompt."""
    try:
        image = pipe(
        prompt,
        height=1024,
        width=1024,
        guidance_scale=3.5,
        num_inference_steps=50,
        max_sequence_length=512,
        generator=torch.Generator("cpu").manual_seed(0)
).images[0]
    except Exception as e:
        # Optionally, handle errors (you can also return a default error image)
        image = None
    return image

# Build the Gradio interface with a two-step process
with gr.Blocks(css=".gradio-container {background-color: #f9f9f9; padding: 20px;}") as demo:
    gr.Markdown("# 2-Step Image Generator")
    gr.Markdown(
        "### Step 1: Prompt Enhancement\n"
        "Enter your original prompt below and click **Enhance Prompt**. "
        "The system will generate a detailed version of your prompt. You can modify the enhanced prompt before generating the image."
    )
    
    with gr.Row():
        original_prompt = gr.Textbox(
            label="Your Original Prompt",
            placeholder="Describe your idea here...",
            lines=3
        )
        enhance_button = gr.Button("Enhance Prompt")
    
    enhanced_prompt_box = gr.Textbox(
        label="Enhanced Prompt (Editable)",
        placeholder="The enhanced prompt will appear here...",
        lines=3
    )
    
    enhance_button.click(fn=enhance_prompt, inputs=original_prompt, outputs=enhanced_prompt_box)
    
    gr.Markdown("### Step 2: Image Generation\n"
                "Review and modify the enhanced prompt if necessary. Once you're ready, click **Generate Image** to create your image.")
    
    generate_button = gr.Button("Generate Image")
    image_output = gr.Image(label="Generated Image")
    
    generate_button.click(fn=generate_image, inputs=enhanced_prompt_box, outputs=image_output)

if __name__ == "__main__":
    demo.launch()