app.py

# --- Filename: app.py ---

import gradio as gr
import openai
import torch
from diffusers import StableDiffusionPipeline, LCMScheduler
import os
from PIL import Image
import io # Required for handling audio file object for OpenAI API
import time # To estimate generation time

# --- Configuration ---
# Load API keys from Hugging Face Secrets or environment variables
# IMPORTANT: Ensure the secret/variable named OPENAI_API_KEY is set in your environment.
openai.api_key = os.environ.get("OPENAI_API_KEY")
hf_token = os.environ.get("HF_TOKEN") # May be needed for model download

if not openai.api_key:
    print("\n" + "="*40)
    print("ERROR: OPENAI_API_KEY environment variable not found.")
    print("Please set the OPENAI_API_KEY secret/variable.")
    print("OpenAI features (prompt enhancement, voice input) WILL FAIL.")
    print("="*40 + "\n")
    # Optionally raise an error or exit if the key is absolutely critical
    # raise ValueError("OpenAI API Key not found!")
else:
    print("OpenAI API Key found.")


# Model IDs
llm_model = "gpt-3.5-turbo"
sd_model_id = "runwayml/stable-diffusion-v1-5"
lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5" # LCM LoRA for faster inference

# Check for GPU availability - WILL BE 'cpu' in your case
device = "cuda" if torch.cuda.is_available() else "cpu"
# Use float32 for CPU for stability/compatibility
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

print(f"Selected Device: {device.upper()}")
print(f"Selected PyTorch Dtype: {torch_dtype}")

# --- Model Loading ---
pipe = None # Initialize pipe to None
try:
    print("Loading Stable Diffusion model... (This might take a while on CPU)")
    pipe = StableDiffusionPipeline.from_pretrained(
        sd_model_id,
        torch_dtype=torch_dtype,
        # use_auth_token=hf_token # Uncomment if you face download issues
    )
    print("Base model loaded. Loading LCM Scheduler and LoRA...")
    # Using LCM Scheduler and LoRA for faster generation
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    pipe.load_lora_weights(lcm_lora_id)
    pipe.fuse_lora() # Fuse LoRA for slightly faster inference after loading
    pipe.to(device) # Move pipe to CPU
    print("Stable Diffusion model loaded successfully with LCM-LoRA on CPU.")
    # Perform a small dummy inference run to warm up / check for errors
    print("Performing a quick warm-up inference...")
    _ = pipe(prompt="warmup", num_inference_steps=1, guidance_scale=1.0, output_type="pil").images[0]
    print("Warm-up successful.")

except Exception as e:
    print(f"\n{'='*40}\nERROR loading Stable Diffusion model: {e}\n{'='*40}\n")
    # pipe remains None, generation will fail gracefully later


# --- Core Functions ---

def enhance_prompt_openai(short_prompt, add_style_keywords):
    """Uses OpenAI LLM to enhance the short prompt."""
    if not openai.api_key:
        # Should not happen if checked at start, but good practice
        return "Error: OpenAI API Key not configured."

    system_message = """You are an expert prompt engineer for text-to-image models like Stable Diffusion.
    Expand the user's short idea into a detailed, vivid, and structured prompt optimized for Stable Diffusion v1.5.
    Include details about the subject, scene, style (e.g., photorealistic, oil painting, cinematic),
    lighting (e.g., soft light, dramatic lighting), composition (e.g., wide shot, close-up),
    and mood. Add high-quality keywords like 'highly detailed', 'sharp focus', 'masterpiece'.
    Keep the prompt concise and effective, ideally under 100 words.""" # Slightly shorter for clarity

    user_message = f"Short idea: \"{short_prompt}\""
    if add_style_keywords:
        user_message += "\nPlease specifically add artistic and quality keywords like 'cinematic lighting', 'photorealistic', '8k', 'masterpiece', 'professional photography'."

    try:
        response = openai.chat.completions.create(
            model=llm_model,
            messages=[
                {"role": "system", "content": system_message},
                {"role": "user", "content": user_message},
            ],
            temperature=0.7,
            max_tokens=150 # Reduced max tokens slightly
        )
        enhanced_prompt = response.choices[0].message.content.strip()
        # Basic cleanup
        enhanced_prompt = enhanced_prompt.replace("Here's a detailed prompt:", "").strip()
        return enhanced_prompt
    except Exception as e:
        print(f"Error calling OpenAI API for prompt enhancement: {e}")
        # Provide a more user-friendly error message
        return f"Error: Could not enhance prompt using OpenAI. ({e})"

def transcribe_audio_openai(audio_path):
    """Transcribes audio using OpenAI Whisper API."""
    if not audio_path:
        return None
    if not openai.api_key:
        print("Warning: OpenAI API Key not configured. Cannot transcribe audio.")
        return "Error: OpenAI API Key needed for transcription."

    try:
        with open(audio_path, "rb") as audio_file:
            transcript = openai.audio.transcriptions.create(
              model="whisper-1",
              file=audio_file
            )
        return transcript.text
    except Exception as e:
        print(f"Error calling OpenAI Whisper API: {e}")
        return f"Error: Could not transcribe audio using OpenAI. ({e})"

def generate_image_lcm(prompt, guidance_scale, num_inference_steps=8): # Increased steps slightly for potentially better quality on CPU
    """Generates an image using the loaded SD+LCM pipeline on CPU."""
    if pipe is None:
        print("Error: Stable Diffusion pipeline is not available.")
        img = Image.new('RGB', (512, 512), color = (128, 128, 128)) # Grey placeholder
        # Add text to placeholder if possible/easy? For now, just grey.
        return img, "Error: Image generation model failed to load."

    print(f"Starting image generation on CPU with prompt: '{prompt}'")
    print(f"Guidance Scale: {guidance_scale}, Steps: {num_inference_steps}. BE PATIENT, THIS WILL BE SLOW.")

    # LCM performs best with low guidance scale
    effective_guidance = max(1.0, min(guidance_scale, 3.0))
    if effective_guidance != guidance_scale:
        print(f"Adjusted guidance scale to {effective_guidance} (optimal range for LCM).")

    negative_prompt = "blurry, low quality, deformed, ugly, text, words, writing, signature, watermark"

    start_time = time.time()
    try:
        # No torch.autocast(device) needed for CPU float32? Check diffusers docs.
        # inference_mode is still good practice
        with torch.inference_mode():
             image = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                guidance_scale=effective_guidance,
                num_inference_steps=num_inference_steps # LCM needs few steps
             ).images[0]
        end_time = time.time()
        duration = end_time - start_time
        print(f"Image generation successful on CPU in {duration:.2f} seconds.")
        return image, f"Image generated in {duration:.2f}s (CPU)." # Return image and status message
    except Exception as e:
        end_time = time.time()
        duration = end_time - start_time
        print(f"Error during image generation after {duration:.2f} seconds: {e}")
        img = Image.new('RGB', (512, 512), color = (255, 100, 100)) # Red-ish placeholder
        return img, f"Error generating image: {e}"

# --- Main Processing Function ---

def process_input(text_input, audio_input, add_style_keywords, guidance_scale):
    """
    Main function triggered by the Gradio interface.
    Handles text/audio input, enhances prompt, generates image.
    """
    status_updates = []
    final_text_input = ""
    enhanced_prompt = ""
    generated_image = None

    # 1. Determine input source
    if text_input and text_input.strip():
        final_text_input = text_input.strip()
        status_updates.append("Using provided text input.")
    elif audio_input:
        status_updates.append("Processing audio input...")
        transcribed_text = transcribe_audio_openai(audio_input)
        if transcribed_text and not transcribed_text.startswith("Error:"):
            final_text_input = transcribed_text
            status_updates.append(f"Transcribed Audio: \"{final_text_input[:100]}...\"" if len(final_text_input) > 100 else f"Transcribed Audio: \"{final_text_input}\"")
        else:
            status_updates.append(transcribed_text or "Error: Transcription failed.") # Show the error message
            final_text_input = "" # Prevent proceeding if transcription fails
    else:
        status_updates.append("Error: Please provide a text description or record audio.")
        # Return current status, empty prompt, no image
        return "\n".join(status_updates), "", None

    # If no valid input text after checking both sources
    if not final_text_input:
         return "\n".join(status_updates), "", None

    # 2. Enhance Prompt
    status_updates.append("Enhancing prompt using OpenAI...")
    if openai.api_key:
        enhanced_prompt = enhance_prompt_openai(final_text_input, add_style_keywords)
        if enhanced_prompt.startswith("Error:"):
            status_updates.append(enhanced_prompt) # Add error to status
            # Decide if we should proceed with the *original* prompt or stop? Let's stop.
            return "\n".join(status_updates), "", None
        else:
             status_updates.append("Prompt enhanced successfully.")
    else:
        status_updates.append("Warning: OpenAI API Key missing. Using original text as prompt.")
        enhanced_prompt = final_text_input # Use original text if API key missing

    # 3. Generate Image
    status_updates.append(f"Generating image on CPU ({device})... **THIS WILL BE SLOW - PLEASE WAIT**")
    # Update the UI status *before* starting generation
    # yield "\n".join(status_updates), enhanced_prompt, None # Requires making the function a generator

    # Simple update (blocks UI until done):
    img_gen_result, img_status_msg = generate_image_lcm(enhanced_prompt, guidance_scale)
    generated_image = img_gen_result
    if img_status_msg:
        status_updates.append(img_status_msg)

    # 4. Return results
    return "\n".join(status_updates), enhanced_prompt, generated_image


# --- Gradio UI ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Prompt Enhancer & Image Generator 🪄🖼️ (CPU Version)")
    gr.Markdown(
        f"**WARNING:** Running on **CPU ({device.upper()})**. Image generation will be **VERY SLOW** (potentially several minutes). Please be patient after clicking Generate."
        f"\nEnter a short description or record audio. It will be enhanced by `{llm_model}` and an image generated using `{sd_model_id}` + LCM acceleration."
    )

    with gr.Row():
        with gr.Column(scale=1):
            # Input Controls
            text_input = gr.Textbox(
                label="Short Description",
                placeholder="e.g., 'magical treehouse in the sky'",
                lines=2
            )
            audio_input = gr.Audio(
                sources=["microphone"],
                type="filepath", # Get file path for OpenAI API
                label="Or Record Audio Input"
            )
            gr.Markdown("---")
            gr.Markdown("**Generation Options**")
            add_style_keywords = gr.Checkbox(
                label="Add Extra Style Keywords (via LLM)?",
                value=True,
                info="Asks the LLM to add 'photorealistic', '8k', 'cinematic' etc."
            )
            guidance_scale = gr.Slider(
                minimum=1.0,
                maximum=3.0, # Keep low for LCM
                step=0.1,
                value=1.5, # Good default for LCM
                label="Guidance Scale",
                info="How closely the image follows the prompt (1-2 recommended for LCM)."
            )
            submit_button = gr.Button("Generate ✨ (Will be slow!)", variant="primary")

        with gr.Column(scale=2):
            # Output Area
            status_output = gr.Textbox(
                label="Status Log",
                interactive=False,
                lines=4 # More lines for verbose status
            )
            enhanced_prompt_output = gr.Textbox(
                label="✨ Enhanced Prompt (from LLM)",
                interactive=False,
                lines=4
            )
            image_output = gr.Image(
                label="🖼️ Generated Image (CPU)",
                type="pil",
                interactive=False,
                height=512, # Set fixed height if desired
                # width=512
            )

    # Connect UI elements
    submit_button.click(
        fn=process_input,
        inputs=[
            text_input,
            audio_input,
            add_style_keywords,
            guidance_scale
        ],
        outputs=[
            status_output,
            enhanced_prompt_output,
            image_output
        ]
    )

    # Clear inputs upon submission for better UX
    submit_button.click(lambda: ("", None), inputs=[], outputs=[text_input, audio_input])


# --- Launch the App ---
if __name__ == "__main__":
    print("\nLaunching Gradio App...")
    # Enable queue for better handling, especially with slow generation
    # share=True can create a public link if running locally (use with caution)
    demo.queue().launch(debug=False, share=False)