Upload 10 files

README.md

@@ -1,12 +1,13 @@
 ---
-title: Paint
-emoji: 🚀
-colorFrom: indigo
-colorTo: pink
+title: Prompt Image
+emoji: 🐨
+colorFrom: blue
+colorTo: yellow
 sdk: gradio
 sdk_version: 5.5.0
 app_file: app.py
 pinned: false
+hf_oauth: false
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

api.py

@@ -0,0 +1,145 @@
+import json
+import logging
+from openai import OpenAI
+from typing import Dict, Any, Optional
+import gradio as gr
+from prompts import PROMPT_ANALYZER_TEMPLATE
+import time
+
+logger = logging.getLogger(__name__)
+
+FALLBACK_MODELS = [
+    "mixtral-8x7b-32768",
+    "llama-3.1-70b-versatile",
+    "llama-3.1-8b-instant",
+    "llama3-70b-8192",
+    "llama3-8b-8192"
+]
+
+class ModelManager:
+    def __init__(self):
+        self.current_model_index = 0
+        self.max_retries = len(FALLBACK_MODELS)
+
+    @property
+    def current_model(self) -> str:
+        return FALLBACK_MODELS[self.current_model_index]
+
+    def next_model(self) -> str:
+        self.current_model_index = (self.current_model_index + 1) % len(FALLBACK_MODELS)
+        logger.info(f"Switching to model: {self.current_model}")
+        return self.current_model
+
+class PromptEnhancementAPI:
+    def __init__(self, api_key: str, base_url: Optional[str] = None):
+        self.client = OpenAI(
+            api_key=api_key,
+            base_url=base_url or "https://api.groq.com/openai/v1"
+        )
+        self.model_manager = ModelManager()
+
+    def _try_parse_json(self, content: str, retries: int = 0) -> Dict[str, Any]:
+        try:
+            result = json.loads(content.strip().lstrip('\n'))
+            if not isinstance(result, dict):
+                raise ValueError("Response is not a valid JSON object")
+            return result
+        except (json.JSONDecodeError, ValueError) as e:
+            if retries < self.model_manager.max_retries - 1:
+                logger.warning(f"JSON parsing failed with model {self.model_manager.current_model}. Switching models...")
+                self.model_manager.next_model()
+                raise e
+            logger.error(f"JSON parsing failed with all models: {str(e)}")
+            raise
+
+    def generate_enhancement(self, system_prompt: str, user_prompt: str, user_directive: str = "", state: Optional[Dict] = None) -> Dict[str, Any]:
+        retries = 0
+        last_error = None
+
+        while retries < self.model_manager.max_retries:
+            try:
+                messages = [
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt}
+                ]
+                
+                if user_directive:
+                    messages.append({"role": "user", "content": f"User directive: {user_directive}"})
+
+                if state:
+                    messages.append({
+                        "role": "assistant",
+                        "content": json.dumps(state)
+                    })
+
+                response = self.client.chat.completions.create(
+                    model=self.model_manager.current_model,
+                    messages=messages,
+                    temperature=0.7,
+                    max_tokens=4000,
+                    response_format={"type": "json_object"}
+                )
+                
+                result = self._try_parse_json(response.choices[0].message.content, retries)
+                return result
+
+            except (json.JSONDecodeError, ValueError) as e:
+                last_error = e
+                retries += 1
+                if retries < self.model_manager.max_retries:
+                    logger.warning(f"Attempt {retries} failed. Switching models and retrying...")
+                    time.sleep(1)  # Brief pause before retry
+                    continue
+                break
+
+            except Exception as e:
+                logger.error(f"API error: {str(e)}")
+                if "rate limit" in str(e).lower():
+                    if retries < self.model_manager.max_retries - 1:
+                        self.model_manager.next_model()
+                        retries += 1
+                        time.sleep(1)
+                        continue
+                raise gr.Error(f"API request failed: {str(e)}")
+
+        logger.error(f"All models failed to generate valid JSON: {str(last_error)}")
+        return create_error_response(user_prompt, user_directive)
+
+class PromptEnhancementSystem:
+    def __init__(self, api_key: str, base_url: Optional[str] = None):
+        self.api = PromptEnhancementAPI(api_key, base_url)
+        self.current_state = None
+        self.history = []
+
+    def start_session(self, prompt: str, user_directive: str = "") -> Dict[str, Any]:
+        formatted_system_prompt = PROMPT_ANALYZER_TEMPLATE.format(
+            input_prompt=prompt,
+            user_directive=user_directive
+        )
+        
+        result = self.api.generate_enhancement(
+            system_prompt=formatted_system_prompt,
+            user_prompt=prompt,
+            user_directive=user_directive
+        )
+        
+        self.current_state = result
+        self.history = [result]
+        return result
+
+    def apply_enhancement(self, choice: str, user_directive: str = "") -> Dict[str, Any]:
+        formatted_system_prompt = PROMPT_ANALYZER_TEMPLATE.format(
+            input_prompt=choice,
+            user_directive=user_directive
+        )
+        
+        result = self.api.generate_enhancement(
+            system_prompt=formatted_system_prompt,
+            user_prompt=choice,
+            user_directive=user_directive,
+            state=self.current_state
+        )
+        
+        self.current_state = result
+        self.history.append(result)
+        return result

app.py

@@ -0,0 +1,30 @@
+import os
+import logging
+from ui import create_interface
+from huggingface_hub import login
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Environment variables check
+required_vars = ["HF_TOKEN", "GROQ_API_KEY"]
+missing_vars = [var for var in required_vars if not os.getenv(var)]
+if missing_vars:
+    raise ValueError(f"Missing required environment variables: {', '.join(missing_vars)}")
+
+# Hugging Face login
+try:
+    login(token=os.getenv("HF_TOKEN"))
+    logger.info("Successfully logged in to Hugging Face")
+except Exception as e:
+    logger.error(f"Failed to login to Hugging Face: {str(e)}")
+    raise
+
+if __name__ == "__main__":
+    try:
+        demo = create_interface()
+        demo.queue(max_size=5).launch()
+    except Exception as e:
+        logger.error(f"Application startup error: {str(e)}")
+        raise

gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

models.py

@@ -0,0 +1,50 @@
+from pydantic import BaseModel, Field, field_validator
+from typing import List, Dict, Any
+
+class ProgressMeters(BaseModel):
+    technical_detail: int = Field(default=0, ge=0, le=100)
+    artistic_style: int = Field(default=0, ge=0, le=100)
+    composition: int = Field(default=0, ge=0, le=100)
+    context: int = Field(default=0, ge=0, le=100)
+
+class SubjectAnalysis(BaseModel):
+    clarity: int = Field(default=0, ge=0, le=100)
+    details_present: List[str] = []
+    details_missing: List[str] = []
+
+class StyleEvaluation(BaseModel):
+    defined_elements: List[str] = []
+    missing_elements: List[str] = []
+    style_score: int = Field(default=0, ge=0, le=100)
+
+class TechnicalAssessment(BaseModel):
+    specified_elements: List[str] = []
+    missing_elements: List[str] = []
+    technical_score: int = Field(default=0, ge=0, le=100)
+
+class CompositionReview(BaseModel):
+    strengths: List[str] = []
+    weaknesses: List[str] = []
+    composition_score: int = Field(default=0, ge=0, le=100)
+
+class InitialAnalysis(BaseModel):
+    subject_analysis: SubjectAnalysis = SubjectAnalysis()
+    style_evaluation: StyleEvaluation = StyleEvaluation()
+    technical_assessment: TechnicalAssessment = TechnicalAssessment()
+    composition_review: CompositionReview = CompositionReview()
+
+class EnhancedVersion(BaseModel):
+    focus_area: str = ""
+    enhanced_prompt: str = ""
+    improvement_score: int = Field(default=0, ge=0, le=100)
+
+class PromptAnalysis(BaseModel):
+    initial_analysis: InitialAnalysis = InitialAnalysis()
+    enhanced_versions: List[EnhancedVersion] = []
+    session_state: Dict[str, Any] = {}
+
+    @field_validator('enhanced_versions', mode='before')
+    def validate_enhanced_versions(cls, v):
+        if not isinstance(v, list):
+            return []
+        return v

prompts.py

@@ -0,0 +1,186 @@
+PROMPT_ANALYZER_TEMPLATE = '''You are a Prompt Enhancement Specialist for image generation. Your task is to analyze a given prompt and dynamically determine the most relevant improvement axes based on the current analysis, while ensuring compliance with specific user directives.
+
+For the following prompt and user directive: 
+<input_prompt>
+{input_prompt}
+</input_prompt>
+
+<user_directive>
+{user_directive}
+</user_directive>
+
+1. Initial Analysis (Comprehensive evaluation of current elements):
+
+Subject Analysis:
+- Main subject identification and clarity
+- Subject details and characteristics
+- Secondary elements and their relationship
+- Scale and proportions
+
+Style Elements:
+- Artistic style presence/absence
+- Medium specification
+- Art movement references
+- Artist influences
+- Historical or cultural context
+
+Technical Specifications:
+- Lighting details
+- Color palette
+- Texture information
+- Resolution indicators
+- Camera angle/perspective
+- Shot type/framing
+
+Compositional Elements:
+- Spatial arrangement
+- Foreground/background balance
+- Rule of thirds consideration
+- Leading lines
+- Focal point clarity
+
+Environmental Context:
+- Setting details
+- Time period
+- Weather/atmospheric conditions
+- Environmental interaction
+- Scene depth
+
+Mood and Atmosphere:
+- Emotional tone
+- Atmospheric qualities
+- Dynamic vs static elements
+- Story/narrative elements
+- Symbolic elements
+
+2. Limitations Assessment:
+- Missing critical details
+- Ambiguous elements
+- Technical omissions
+- Stylistic gaps
+- Compositional weaknesses
+- Context deficiencies
+- Mood/atmosphere undefined areas
+
+3. Improvement Axes (Select 4 most impactful):
+For each axis, consider:
+- Impact on visual outcome
+- Technical feasibility
+- AI model capabilities
+- Balance between specificity and creativity
+- Enhancement of original vision
+- Visual interest addition
+- Technical precision improvement
+- User directive compliance and integration
+- ...
+
+4. Enhancement Strategy:
+For each improvement axis:
+- Specific terminology to add
+- Technical parameters to include
+- Stylistic elements to incorporate
+- Compositional guidance
+- Atmospheric elements
+- Reference points (artists, styles, techniques)
+- User directive implementation methods
+
+Now provide your analysis in this JSON structure:
+
+{{
+    "initial_analysis": {{
+        "initial_prompt": {input_prompt},
+        "user_directive": {user_directive},
+        "directive_impact_assessment": {{
+            "feasibility": string,
+            "integration_approach": string,
+            "potential_conflicts": [string],
+            "resolution_strategy": string
+        }},
+        "subject_analysis": {{
+            "score": integer(0-100),
+            "strengths": [string],
+            "weaknesses": [string]
+        }},
+        "style_evaluation": {{
+            "score": integer(0-100),
+            "strengths": [string],
+            "weaknesses": [string]
+        }},
+        "technical_assessment": {{
+            "score": integer(0-100),
+            "strengths": [string],
+            "weaknesses": [string]
+        }},
+        "composition_review": {{
+            "score": integer(0-100),
+            "strengths": [string],
+            "weaknesses": [string]
+        }},
+        "context_evaluation": {{
+            "score": integer(0-100),
+            "strengths": [string],
+            "weaknesses": [string]
+        }},
+        "mood_assessment": {{
+            "score": integer(0-100),
+            "strengths": [string],
+            "weaknesses": [string]
+        }}
+    }},
+    "improvement_axes": [
+        {{
+            "axis_name": string,
+            "focus_area": string,
+            "version": integer,
+            "score": integer(0-100),
+            "current_state": string,
+            "directive_alignment": string,
+            "recommended_additions": [string],
+            "expected_impact": string,
+            "technical_considerations": [string],
+            "enhanced_prompt": string,
+            "expected_improvements": [string]
+        }}
+    ],
+    "technical_recommendations": {{
+        "style_keywords": [string],
+        "composition_tips": [string],
+        "negative_prompt_suggestions": [string],
+        "directive_specific_adjustments": [string]
+    }}
+}}
+
+Guidelines for Dynamic Enhancement:
+1. Analyze current scores to identify weakest areas
+2. Ensure all improvements align with the user directive (if provided)
+3. Consider improvement potential for each axis
+4. Select 4 most impactful axes based on:
+   - User directive compliance (highest priority if provided)
+   - Current analysis scores
+   - Previous improvements
+   - Remaining potential
+   - Overall image quality goals
+5. Generate targeted enhancements for selected axes
+
+Remember to:
+- Prioritize user directive implementation while maintaining prompt integrity
+- Keep improvements relevant to image generation
+- Maintain the original intent of the prompt
+- Be specific and detailed in suggestions
+- Ensure each enhanced version builds on the original
+- Focus on visual elements that AI image generators understand
+- Consider technical aspects like lighting, composition, and style
+- Add specific artistic references when relevant
+- Balance detail with creativity
+- Consider AI model capabilities and limitations
+- Provide practical composition guidance
+- Include relevant style keywords
+- Specify negative prompt elements
+
+Each iteration should:
+1. Verify user directive compliance
+2. Reassess current state
+3. Identify new priority areas
+4. Generate fresh improvement approaches
+5. Build upon previous enhancements while maintaining user directive alignment
+6. Maintain coherence with original concept'''

requirements.txt

@@ -0,0 +1,11 @@
+accelerate
+git+https://github.com/huggingface/diffusers.git
+invisible_watermark
+torch
+transformers==4.42.4
+xformers
+sentencepiece
+gradio==4.14.0
+numpy==1.24.3
+openai==1.3.0
+huggingface-hub>=0.19.0

schemas.py

@@ -0,0 +1,79 @@
+from typing import List, Optional, Dict, Any
+from pydantic import BaseModel, Field, ConfigDict
+
+class DirectiveImpactAssessment(BaseModel):
+    feasibility: str = Field(default="Not assessed")
+    integration_approach: str = Field(default="Not determined")
+    potential_conflicts: List[str] = Field(default_factory=lambda: ["None identified"])
+    resolution_strategy: str = Field(default="Not required")
+
+class AnalysisScore(BaseModel):
+    score: int = Field(default=0, ge=0, le=100)
+    strengths: List[str] = Field(default_factory=lambda: ["Not analyzed"])
+    weaknesses: List[str] = Field(default_factory=lambda: ["Not analyzed"])
+
+class ImprovementAxis(BaseModel):
+    axis_name: str = Field(default="Default")
+    focus_area: str = Field(default="Not specified")
+    version: int = Field(default=1)
+    score: int = Field(default=0, ge=0, le=100)
+    current_state: str = Field(default="Not evaluated")
+    directive_alignment: str = Field(default="Not aligned")
+    recommended_additions: List[str] = Field(default_factory=lambda: ["No recommendations"])
+    expected_impact: str = Field(default="Not determined")
+    technical_considerations: List[str] = Field(default_factory=lambda: ["None specified"])
+    enhanced_prompt: str = Field(default="")
+    expected_improvements: List[str] = Field(default_factory=lambda: ["None specified"])
+
+class TechnicalRecommendations(BaseModel):
+    style_keywords: List[str] = Field(default_factory=lambda: ["None"])
+    composition_tips: List[str] = Field(default_factory=lambda: ["None"])
+    negative_prompt_suggestions: List[str] = Field(default_factory=lambda: ["None"])
+    directive_specific_adjustments: List[str] = Field(default_factory=lambda: ["None"])
+
+class InitialAnalysis(BaseModel):
+    initial_prompt: str
+    user_directive: str = Field(default="")
+    directive_impact_assessment: DirectiveImpactAssessment = Field(default_factory=DirectiveImpactAssessment)
+    subject_analysis: AnalysisScore = Field(default_factory=AnalysisScore)
+    style_evaluation: AnalysisScore = Field(default_factory=AnalysisScore)
+    technical_assessment: AnalysisScore = Field(default_factory=AnalysisScore)
+    composition_review: AnalysisScore = Field(default_factory=AnalysisScore)
+    context_evaluation: AnalysisScore = Field(default_factory=AnalysisScore)
+    mood_assessment: AnalysisScore = Field(default_factory=AnalysisScore)
+
+class APIResponse(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    initial_analysis: InitialAnalysis
+    improvement_axes: List[ImprovementAxis] = Field(default_factory=list)
+    technical_recommendations: TechnicalRecommendations = Field(default_factory=TechnicalRecommendations)
+
+def create_error_response(user_prompt: str, user_directive: str = "") -> Dict[str, Any]:
+    """Create a standardized error response that complies with APIResponse model"""
+    return APIResponse(
+        initial_analysis=InitialAnalysis(
+            initial_prompt=user_prompt,
+            user_directive=user_directive
+        ),
+        improvement_axes=[
+            ImprovementAxis(
+                axis_name="Error",
+                focus_area="Error occurred",
+                version=1,
+                score=0,
+                current_state="Failed",
+                directive_alignment="Failed to assess",
+                recommended_additions=["Error processing prompt"],
+                expected_impact="None",
+                technical_considerations=["Error occurred"],
+                enhanced_prompt=user_prompt,
+                expected_improvements=["Error processing prompt"]
+            )
+        ],
+        technical_recommendations=TechnicalRecommendations(
+            style_keywords=["Error"],
+            composition_tips=["Error"],
+            negative_prompt_suggestions=["Error"],
+            directive_specific_adjustments=["Error"]
+        )
+    ).model_dump()

ui.py

@@ -0,0 +1,480 @@
+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,
+            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()

ui_old.py

@@ -0,0 +1,346 @@
+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,
+    seed=42,
+    randomize_seed=False,
+    width=512,
+    height=512,
+    num_inference_steps=4,
+    progress=gr.Progress(track_tqdm=True)
+):
+    try:
+        # Extract prompts from improvement axes
+        prompts = [axis["enhanced_prompt"] for axis in improvement_axes if axis.get("enhanced_prompt")]
+        
+        if not prompts:
+            return [None] * 4 + [seed]
+        
+        if randomize_seed:
+            current_seed = random.randint(0, MAX_SEED)
+        else:
+            current_seed = seed
+            
+        print(f"Generating images with {len(prompts)} prompts")
+        print(f"Using seed: {current_seed}")
+        
+        # Generate all images in a single batch
+        generator = torch.Generator().manual_seed(current_seed)
+        images = pipe(
+            prompt=prompts,  # Pass list of prompts directly
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            guidance_scale=0.0
+        ).images
+        
+        # Pad with None if we have fewer than 4 images
+        while len(images) < 4:
+            images.append(None)
+            
+        print("All images generated successfully")
+        return images[:4] + [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):
+    """Handle image selection event"""
+    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 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):
+        try:
+            print(f"\n=== Processing prompt: {prompt}")
+            state = system.start_session(prompt)
+            
+            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", {}),
+                None, None, None, None,  # Four None values for the four image outputs
+                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 + [{}, {}, None, None, None, None, {}] + [gr.update(visible=False)] * 4
+
+    def handle_option_click(option_num, input_prompt, current_text):
+        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.Row():
+            input_prompt = gr.Textbox(
+                label="Initial Prompt",
+                placeholder="Enter your prompt here...",
+                lines=3,
+                scale=1
+            )
+            current_prompt = gr.Textbox(
+                label="Current Prompt",
+                lines=3,
+                scale=1,
+                interactive=True
+            )
+
+        with gr.Row():
+            start_btn = gr.Button("Start Enhancement", variant="primary")
+
+        with gr.Row():
+            option_buttons = [gr.Button("", visible=False) for _ in range(4)]
+
+        with gr.Tabs():
+            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.TabItem("Generated Images"):
+                with gr.Row():
+                    generated_images = [
+                        gr.Image(
+                            label=f"Image {i+1}",
+                            type="pil",
+                            show_label=True,
+                            height=256,
+                            width=256,
+                            interactive=True,
+                            elem_id=f"image_{i}"
+                        ) for i in range(4)
+                    ]
+                
+                with gr.Row():
+                    finalize_btn = gr.Button("Generate All Images", variant="primary")
+            
+
+                with gr.Accordion("Image Generation Settings", open=False):
+                    with gr.Row():
+                        seed = gr.Slider(
+                            label="Seed",
+                            minimum=0,
+                            maximum=2048,
+                            step=1,
+                            value=42
+                        )
+                        randomize_seed = gr.Checkbox(
+                            label="Randomize seed",
+                            value=True
+                        )
+                    with gr.Row():
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=2048,
+                            step=256,
+                            value=512
+                        )
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=2048,
+                            step=256,
+                            value=512
+                        )
+                        num_inference_steps = gr.Slider(
+                            label="Steps",
+                            minimum=1,
+                            maximum=50,
+                            step=1,
+                            value=4
+                        )
+
+        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 select events for each image
+        for i, img in enumerate(generated_images):
+            img.select(
+                fn=handle_image_select,
+                inputs=[improvement_axes],
+                outputs=[input_prompt]
+            )
+
+        start_btn.click(
+            update_interface,
+            inputs=[input_prompt],
+            outputs=[
+                input_prompt,
+                current_prompt,
+                subject_analysis,
+                style_evaluation,
+                technical_assessment,
+                composition_review,
+                context_evaluation,
+                mood_assessment,
+                improvement_axes,
+                technical_recommendations
+            ] + generated_images + [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
+                ],
+                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,
+                seed,
+                randomize_seed,
+                width,
+                height,
+                num_inference_steps
+            ],
+            outputs=generated_images + [seed]
+        )
+
+    print("Interface setup complete")
+    return interface