ui.py

import spaces
import os
import gradio as gr
import random
import torch
import logging
import numpy as np
from typing import Dict, Any, List
from diffusers import DiffusionPipeline
from api import PromptEnhancementSystem

# Constants
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODEL_ID = "black-forest-labs/FLUX.1-schnell"
DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32

print(f"Using device: {DEVICE}")
logger = logging.getLogger(__name__)

# Initialize model
try:
    print("Loading model...")
    pipe = DiffusionPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=DTYPE
    ).to(DEVICE)
    print("Model loaded successfully")
    logger.info("Model loaded successfully")
except Exception as e:
    print(f"Failed to load model: {str(e)}")
    logger.error(f"Failed to load model: {str(e)}")
    raise

@spaces.GPU()
def generate_multiple_images_batch(
    improvement_axes,
    current_gallery,
    seed=42,
    randomize_seed=False,
    width=512,
    height=512,
    num_inference_steps=4,
    current_prompt="",
    initial_prompt="",
    progress=gr.Progress(track_tqdm=True)
):
    try:
        # Use current_prompt if not empty, otherwise fall back to initial_prompt
        input_prompt = current_prompt if current_prompt.strip() else initial_prompt
        
        # Extract prompts from improvement axes or use the input prompt if no axes
        prompts = [axis["enhanced_prompt"] for axis in improvement_axes if axis.get("enhanced_prompt")]
        if not prompts and input_prompt:
            prompts = [input_prompt]
            
        if not prompts:
            return [None] * 4 + [current_gallery] + [seed]

        if randomize_seed:
            current_seed = random.randint(0, MAX_SEED)
        else:
            current_seed = seed

        print(f"Generating images with prompt: {input_prompt}")
        print(f"Using seed: {current_seed}")

        # Generate images with the selected prompt
        generator = torch.Generator().manual_seed(current_seed)
        images = pipe(
            prompt=prompts,
            width=width,
            height=height,
            num_inference_steps=num_inference_steps,
            generator=generator,
            max_sequence_length=256,  # Maximum allowed for schnell
            guidance_scale=0.0
        ).images

        # Pad with None if we have fewer than 4 images
        while len(images) < 4:
            images.append(None)

        # Update gallery with new images
        current_gallery = current_gallery or []
        new_gallery = current_gallery + [(img, f"Prompt: {prompt}") for img, prompt in zip(images, prompts) if img is not None]

        print("All images generated successfully")
        return images[:4] + [new_gallery] + [current_seed]

    except Exception as e:
        print(f"Image generation error: {str(e)}")
        logger.error(f"Image generation error: {str(e)}")
        raise

def handle_image_select(evt: gr.SelectData, improvement_axes_data):
    try:
        if improvement_axes_data and isinstance(improvement_axes_data, list):
            selected_index = evt.index[1] if isinstance(evt.index, tuple) else evt.index
            if selected_index < len(improvement_axes_data):
                selected_prompt = improvement_axes_data[selected_index].get("enhanced_prompt", "")
                return selected_prompt
        return ""
    except Exception as e:
        print(f"Error in handle_image_select: {str(e)}")
        return ""

def handle_gallery_select(evt: gr.SelectData, gallery_data):
    try:
        if gallery_data and isinstance(evt.index, int) and evt.index < len(gallery_data):
            image, prompt = gallery_data[evt.index]
            # Remove "Prompt: " prefix if it exists
            prompt = prompt.replace("Prompt: ", "") if prompt else ""
            return {"prompt": prompt}, prompt
        return None, ""
    except Exception as e:
        print(f"Error in handle_gallery_select: {str(e)}")
        return None, ""
        
def clear_gallery():
    return [], None, None, None, None  # Returns empty gallery and clears the 4 images

def zip_gallery_images(gallery):
    try:
        if not gallery:
            return None
            
        import io
        import zipfile
        from datetime import datetime
        import numpy as np
        from PIL import Image
        
        # Create zip file in memory
        zip_buffer = io.BytesIO()
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"gallery_images_{timestamp}.zip"
        
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
            for i, (img_data, prompt) in enumerate(gallery):
                try:
                    if img_data is not None:
                        # Convert numpy array to PIL Image if needed
                        if isinstance(img_data, np.ndarray):
                            img = Image.fromarray(np.uint8(img_data))
                        elif isinstance(img_data, Image.Image):
                            img = img_data
                        else:
                            print(f"Skipping image {i}: invalid type {type(img_data)}")
                            continue
                            
                        # Save image to bytes
                        img_buffer = io.BytesIO()
                        img.save(img_buffer, format='PNG')
                        img_buffer.seek(0)
                        
                        # Create filename with prompt
                        safe_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '-', '_')).strip()
                        img_filename = f"image_{i+1}_{safe_prompt}.png"
                        
                        # Add to zip
                        zip_file.writestr(img_filename, img_buffer.getvalue())
                except Exception as img_error:
                    print(f"Error processing image {i}: {str(img_error)}")
                    continue
        
        # Prepare zip for download
        zip_buffer.seek(0)
        
        # Return the file data and name
        return {
            "name": filename,
            "data": zip_buffer.getvalue()
        }
        
    except Exception as e:
        print(f"Error creating zip: {str(e)}")
        return None
        

def create_interface():
    print("Creating interface...")
    api_key = os.getenv("GROQ_API_KEY")
    base_url = os.getenv("API_BASE_URL")

    if not api_key:
        print("GROQ_API_KEY not found in environment variables")
        raise ValueError("GROQ_API_KEY not found in environment variables")

    system = PromptEnhancementSystem(api_key, base_url)
    print("PromptEnhancementSystem initialized")

    def update_interface(prompt, user_directive):
        try:
            print(f"\n=== Processing prompt: {prompt}")
            print(f"User directive: {user_directive}")
            state = system.start_session(prompt, user_directive)
            improvement_axes = state.get("improvement_axes", [])
            initial_analysis = state.get("initial_analysis", {})
            enhanced_prompt = ""
            if improvement_axes and len(improvement_axes) > 0:
                enhanced_prompt = improvement_axes[0].get("enhanced_prompt", prompt)

            button_updates = []
            for i in range(4):
                if i < len(improvement_axes):
                    focus_area = improvement_axes[i].get("focus_area", f"Option {i+1}")
                    button_updates.append(gr.update(visible=True, value=focus_area))
                else:
                    button_updates.append(gr.update(visible=False))

            return [prompt, enhanced_prompt] + [
                initial_analysis.get(key, {}) for key in [
                    "subject_analysis",
                    "style_evaluation",
                    "technical_assessment",
                    "composition_review",
                    "context_evaluation",
                    "mood_assessment"
                ]
            ] + [
                improvement_axes,
                state.get("technical_recommendations", {}),
                state
            ] + button_updates

        except Exception as e:
            print(f"Error in update_interface: {str(e)}")
            logger.error(f"Error in update_interface: {str(e)}")
            empty_analysis = {"score": 0, "strengths": [], "weaknesses": ["Error occurred"]}
            return [prompt, prompt] + [empty_analysis] * 6 + [{}, {}, {}] + [gr.update(visible=False)] * 4

    def handle_option_click(option_num, input_prompt, current_text, user_directive):
        try:
            print(f"\n=== Processing option {option_num}")
            state = system.current_state
            if state and "improvement_axes" in state:
                improvement_axes = state["improvement_axes"]
                if option_num < len(improvement_axes):
                    selected_prompt = improvement_axes[option_num]["enhanced_prompt"]
                    return [
                        input_prompt,
                        selected_prompt,
                        state.get("initial_analysis", {}).get("subject_analysis", {}),
                        state.get("initial_analysis", {}).get("style_evaluation", {}),
                        state.get("initial_analysis", {}).get("technical_assessment", {}),
                        state.get("initial_analysis", {}).get("composition_review", {}),
                        state.get("initial_analysis", {}).get("context_evaluation", {}),
                        state.get("initial_analysis", {}).get("mood_assessment", {}),
                        improvement_axes,
                        state.get("technical_recommendations", {}),
                        state
                    ]
            return handle_error()
        except Exception as e:
            print(f"Error in handle_option_click: {str(e)}")
            logger.error(f"Error in handle_option_click: {str(e)}")
            return handle_error()

    def handle_error():
        empty_analysis = {"score": 0, "strengths": [], "weaknesses": ["Error occurred"]}
        return ["", "", empty_analysis, empty_analysis, empty_analysis, empty_analysis, empty_analysis, empty_analysis, [], {}, {}]

    with gr.Blocks(
        title="AI Prompt Enhancement System",
        theme=gr.themes.Soft(),
        css="footer {visibility: hidden}"
    ) as interface:
        gr.Markdown("# 🎨 AI Prompt Enhancement & Image Generation System")

        with gr.TabItem("Images Generation"):
            with gr.Row():
                input_prompt = gr.Textbox(
                    label="Initial Prompt",
                    placeholder="Enter your prompt here...",
                    lines=3,
                    scale=1
                )
    
            with gr.Row():
                user_directive = gr.Textbox(
                    label="User Directive",
                    placeholder="Enter specific requirements...",
                    lines=2,
                    scale=1
                )

            with gr.Row():
                start_btn = gr.Button("Start Enhancement", variant="primary")
            with gr.Row():
                current_prompt = gr.Textbox(
                    label="Current Prompt",
                    lines=3,
                    scale=1,
                    interactive=True
                )
            with gr.Row():
                option_buttons = [gr.Button("", visible=False) for _ in range(4)]    
            with gr.Row():
                finalize_btn = gr.Button("Generate Images", variant="primary")
            with gr.Row():
                generated_images = [
                    gr.Image(
                        label=f"Image {i+1}",
                        type="pil",
                        show_label=False,
                        height=256,
                        width=256,
                        interactive=False,
                        show_download_button=False,
                        elem_id=f"image_{i}"
                    ) for i in range(4)
                ]
                        
        with gr.TabItem("Images Gallery"):
            with gr.Row():
                image_gallery = gr.Gallery(
                    label="Generated Images History",
                    show_label=False,
                    columns=4,
                    rows=None,
                    height=800,
                    object_fit="contain"
                )
            with gr.Row():
                clear_gallery_btn = gr.Button("Clear Gallery", variant="secondary")
            with gr.Row():
                selected_image_data = gr.JSON(label="Selected Image Data", visible=True)
                copy_to_prompt_btn = gr.Button("Copy Prompt to Current", visible=True)
        with gr.TabItem("Image Generation Settings"):
            with gr.Row():
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=42
                )
                randomize_seed = gr.Checkbox(
                    label="Randomize seed",
                    value=True
                )

            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=256,
                    value=512
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=256,
                    value=512
                )
                num_inference_steps = gr.Slider(
                    label="Steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4
                )                
        with gr.TabItem("Initial Analysis"):
            with gr.Row():
                with gr.Column():
                    subject_analysis = gr.JSON(label="Subject Analysis")
                with gr.Column():
                    style_evaluation = gr.JSON(label="Style Evaluation")
                with gr.Column():
                    technical_assessment = gr.JSON(label="Technical Assessment")

            with gr.Row():
                with gr.Column():
                    composition_review = gr.JSON(label="Composition Review")
                with gr.Column():
                    context_evaluation = gr.JSON(label="Context Evaluation")
                with gr.Column():
                    mood_assessment = gr.JSON(label="Mood Assessment")

            with gr.Accordion("Additional Information", open=False):
                improvement_axes = gr.JSON(label="Improvement Axes")
                technical_recommendations = gr.JSON(label="Technical Recommendations")
                full_llm_response = gr.JSON(label="Full LLM Response")

        # Add event handlers
        for i, img in enumerate(generated_images):
            img.select(
                fn=handle_image_select,
                inputs=[improvement_axes],
                outputs=[current_prompt],
                show_progress=False
            )

        start_btn.click(
            update_interface,
            inputs=[input_prompt, user_directive],
            outputs=[
                input_prompt,
                current_prompt,
                subject_analysis,
                style_evaluation,
                technical_assessment,
                composition_review,
                context_evaluation,
                mood_assessment,
                improvement_axes,
                technical_recommendations,
                full_llm_response
            ] + option_buttons
        )

        for i, btn in enumerate(option_buttons):
            btn.click(
                handle_option_click,
                inputs=[
                    gr.Slider(value=i, visible=False),
                    input_prompt,
                    current_prompt,
                    user_directive
                ],
                outputs=[
                    input_prompt,
                    current_prompt,
                    subject_analysis,
                    style_evaluation,
                    technical_assessment,
                    composition_review,
                    context_evaluation,
                    mood_assessment,
                    improvement_axes,
                    technical_recommendations,
                    full_llm_response
                ]
            )

        finalize_btn.click(
            generate_multiple_images_batch,
            inputs=[
                improvement_axes,
                image_gallery,
                seed,
                randomize_seed,
                width,
                height,
                num_inference_steps,
                current_prompt,
                input_prompt
            ],
            outputs=generated_images + [image_gallery] + [seed]
        )

        clear_gallery_btn.click(
            clear_gallery,
            inputs=[],
            outputs=[image_gallery] + generated_images
        )
        
        # Add gallery selection handler
        image_gallery.select(
            fn=handle_gallery_select,
            inputs=[image_gallery],
            outputs=[selected_image_data, current_prompt]
        )
        
        # Add copy button handler
        # Fix the copy button handler by adding a null check
        copy_to_prompt_btn.click(
            lambda x: x["prompt"] if x and isinstance(x, dict) and "prompt" in x else "",
            inputs=[selected_image_data],
            outputs=[current_prompt]
        )
        print("Interface setup complete")
        return interface

if __name__ == "__main__":
    interface = create_interface()
    interface.launch()