Update app.py

app.py

@@ -5,6 +5,7 @@ import os
 import time
 import torch
 from diffusers import FluxPipeline
+from transformers import pipeline
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {DEVICE}")
@@ -16,8 +17,103 @@ DEFAULT_NUM_INFERENCE_STEPS = 15
 DEFAULT_MAX_SEQUENCE_LENGTH = 512
 HF_TOKEN = os.environ.get("HF_ACCESS_TOKEN")
 
-# Cache for the pipeline
+# Cache for the pipelines
 CACHED_PIPE = None
+CACHED_LLM_PIPE = None
+
+def load_llm_pipeline():
+    """Load the LLM pipeline for prompt enhancement"""
+    global CACHED_LLM_PIPE
+    if CACHED_LLM_PIPE is not None:
+        return CACHED_LLM_PIPE
+    
+    print("Loading LLM pipeline for prompt enhancement...")
+    try:
+        # Note: Using a smaller model that's actually available
+        # You can replace this with "openai/gpt-oss-120b" if you have access
+        llm_pipe = pipeline(
+            "text-generation", 
+            model="microsoft/Phi-3-mini-4k-instruct",  # Alternative smaller model
+            torch_dtype=torch.bfloat16,
+            device_map="auto"
+        )
+        CACHED_LLM_PIPE = llm_pipe
+        print("LLM pipeline loaded successfully")
+        return llm_pipe
+    except Exception as e:
+        print(f"Error loading LLM pipeline: {e}")
+        # Fallback to a simpler model if the main one fails
+        try:
+            llm_pipe = pipeline(
+                "text-generation",
+                model="gpt2",  # Fallback to GPT-2
+                device_map="auto"
+            )
+            CACHED_LLM_PIPE = llm_pipe
+            print("Loaded fallback LLM pipeline (GPT-2)")
+            return llm_pipe
+        except Exception as e2:
+            print(f"Error loading fallback LLM pipeline: {e2}")
+            return None
+
+def enhance_prompt(prompt, progress=gr.Progress()):
+    """Enhance the prompt using LLM"""
+    if not prompt:
+        return prompt, "Please enter a prompt first."
+    
+    progress(0.3, desc="Enhancing prompt with AI...")
+    
+    try:
+        llm_pipe = load_llm_pipeline()
+        if llm_pipe is None:
+            return prompt, "LLM pipeline not available, using original prompt."
+        
+        # Create enhancement prompt
+        messages = [
+            {
+                "role": "system", 
+                "content": "You are a helpful assistant that enhances image generation prompts. Make prompts more detailed, artistic, and visually descriptive while keeping the core concept. Add details about lighting, style, colors, mood, and composition. Keep the enhanced prompt under 200 words."
+            },
+            {
+                "role": "user", 
+                "content": f"Enhance this image generation prompt, making it more detailed and artistic: '{prompt}'"
+            }
+        ]
+        
+        # Generate enhanced prompt
+        result = llm_pipe(
+            messages,
+            max_new_tokens=200,
+            temperature=0.7,
+            do_sample=True,
+            top_p=0.9
+        )
+        
+        # Extract the enhanced prompt from the response
+        if isinstance(result, list) and len(result) > 0:
+            enhanced = result[0].get('generated_text', '')
+            # Extract only the assistant's response
+            if isinstance(enhanced, list):
+                for msg in enhanced:
+                    if msg.get('role') == 'assistant':
+                        enhanced = msg.get('content', prompt)
+                        break
+            elif isinstance(enhanced, str):
+                # Clean up the response if needed
+                enhanced = enhanced.strip()
+                if enhanced.startswith("Enhanced prompt:"):
+                    enhanced = enhanced.replace("Enhanced prompt:", "").strip()
+            
+            if enhanced and enhanced != prompt:
+                return enhanced, "Prompt enhanced successfully!"
+            else:
+                return prompt, "Using original prompt."
+        else:
+            return prompt, "Enhancement failed, using original prompt."
+            
+    except Exception as e:
+        print(f"Error during prompt enhancement: {e}")
+        return prompt, f"Enhancement error: {e}. Using original prompt."
 
 def load_bnb_4bit_pipeline():
     """Load the 4-bit quantized pipeline"""
@@ -45,12 +141,19 @@ def load_bnb_4bit_pipeline():
         raise
 
 @spaces.GPU(duration=240)
-def generate_image(prompt, progress=gr.Progress(track_tqdm=True)):
-    """Generate image using 4-bit quantized model"""
+def generate_image(prompt, use_enhancement=False, progress=gr.Progress(track_tqdm=True)):
+    """Generate image using 4-bit quantized model with optional prompt enhancement"""
     if not prompt:
-        return None, "Please enter a prompt."
+        return None, prompt, "Please enter a prompt."
     
-    progress(0.2, desc="Loading 4-bit quantized model...")
+    enhanced_prompt = prompt
+    enhancement_status = ""
+    
+    # Enhance prompt if requested
+    if use_enhancement:
+        enhanced_prompt, enhancement_status = enhance_prompt(prompt, progress)
+    
+    progress(0.5, desc="Loading 4-bit quantized model...")
     
     try:
         # Load the 4-bit pipeline
@@ -58,7 +161,7 @@ def generate_image(prompt, progress=gr.Progress(track_tqdm=True)):
         
         # Set up generation parameters
         pipe_kwargs = {
-            "prompt": prompt,
+            "prompt": enhanced_prompt,
             "height": DEFAULT_HEIGHT,
             "width": DEFAULT_WIDTH,
             "guidance_scale": DEFAULT_GUIDANCE_SCALE,
@@ -70,7 +173,7 @@ def generate_image(prompt, progress=gr.Progress(track_tqdm=True)):
         seed = random.getrandbits(64)
         print(f"Using seed: {seed}")
         
-        progress(0.5, desc="Generating image...")
+        progress(0.7, desc="Generating image...")
         
         # Generate image
         gen_start_time = time.time()
@@ -81,19 +184,29 @@ def generate_image(prompt, progress=gr.Progress(track_tqdm=True)):
         mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0
         print(f"Memory reserved: {mem_reserved:.2f} GB")
         
-        return image, f"Generation complete! (Seed: {seed})"
+        status_msg = f"Generation complete! (Seed: {seed})"
+        if enhancement_status:
+            status_msg = f"{enhancement_status} | {status_msg}"
+        
+        return image, enhanced_prompt, status_msg
         
     except Exception as e:
         print(f"Error during generation: {e}")
-        return None, f"Error: {e}"
+        return None, enhanced_prompt, f"Error: {e}"
+
+@spaces.GPU(duration=60)
+def enhance_only(prompt, progress=gr.Progress()):
+    """Only enhance the prompt without generating an image"""
+    enhanced_prompt, status = enhance_prompt(prompt, progress)
+    return enhanced_prompt, status
 
 # Create Gradio interface
-with gr.Blocks(title="FLUXllama", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="FLUXllama Enhanced", theme=gr.themes.Soft()) as demo:
     gr.HTML(
         """
         <div style='text-align: center; margin-bottom: 20px;'>
-            <h1>FLUXllama</h1>
-            <p>FLUX.1-dev 4-bit Quantized Version</p>
+            <h1>FLUXllama Enhanced</h1>
+            <p>FLUX.1-dev 4-bit Quantized Version with AI Prompt Enhancement</p>
         </div>
         """
     )
@@ -112,14 +225,31 @@ with gr.Blocks(title="FLUXllama", theme=gr.themes.Soft()) as demo:
         """
     )
     
-    with gr.Row():
+    with gr.Column():
         prompt_input = gr.Textbox(
             label="Enter your prompt",
             placeholder="e.g., A photorealistic portrait of an astronaut on Mars",
-            lines=2,
-            scale=4
+            lines=3
         )
-        generate_button = gr.Button("Generate", variant="primary", scale=1)
+        
+        with gr.Row():
+            enhance_checkbox = gr.Checkbox(
+                label="🎨 Use AI Prompt Enhancement",
+                value=False,
+                info="Automatically enhance your prompt for better results"
+            )
+            enhance_only_button = gr.Button("✨ Enhance Only", variant="secondary", scale=1)
+        
+        enhanced_prompt_display = gr.Textbox(
+            label="Enhanced Prompt (will appear after enhancement)",
+            lines=3,
+            interactive=False,
+            visible=True
+        )
+        
+        with gr.Row():
+            generate_button = gr.Button("🚀 Generate Image", variant="primary", scale=2)
+            generate_enhanced_button = gr.Button("🎨 Enhance & Generate", variant="primary", scale=2)
     
     output_image = gr.Image(
         label="Generated Image (4-bit Quantized)",
@@ -135,16 +265,28 @@ with gr.Blocks(title="FLUXllama", theme=gr.themes.Soft()) as demo:
     
     # Connect components
     generate_button.click(
-        fn=generate_image,
+        fn=lambda p: generate_image(p, use_enhancement=False),
+        inputs=[prompt_input],
+        outputs=[output_image, enhanced_prompt_display, status_text]
+    )
+    
+    generate_enhanced_button.click(
+        fn=lambda p: generate_image(p, use_enhancement=True),
+        inputs=[prompt_input],
+        outputs=[output_image, enhanced_prompt_display, status_text]
+    )
+    
+    enhance_only_button.click(
+        fn=enhance_only,
         inputs=[prompt_input],
-        outputs=[output_image, status_text]
+        outputs=[enhanced_prompt_display, status_text]
     )
     
-    # Enter key to submit
+    # Enter key to submit (with enhancement checkbox consideration)
     prompt_input.submit(
         fn=generate_image,
-        inputs=[prompt_input],
-        outputs=[output_image, status_text]
+        inputs=[prompt_input, enhance_checkbox],
+        outputs=[output_image, enhanced_prompt_display, status_text]
     )
     
     # Example prompts
@@ -161,6 +303,26 @@ with gr.Blocks(title="FLUXllama", theme=gr.themes.Soft()) as demo:
         ],
         inputs=prompt_input
     )
+    
+    gr.HTML(
+        """
+        <div style='text-align: center; margin-top: 20px; padding: 20px; background-color: #f0f0f0; border-radius: 10px;'>
+            <h3>✨ Prompt Enhancement Feature</h3>
+            <p>This app now includes AI-powered prompt enhancement! The enhancement feature will:</p>
+            <ul style='text-align: left; display: inline-block;'>
+                <li>Add artistic details and visual descriptions</li>
+                <li>Specify lighting, mood, and atmosphere</li>
+                <li>Include style and composition elements</li>
+                <li>Make your prompts more effective for image generation</li>
+            </ul>
+            <p><strong>How to use:</strong></p>
+            <p>1. Enter a simple prompt</p>
+            <p>2. Click "✨ Enhance Only" to preview the enhanced version</p>
+            <p>3. Click "🎨 Enhance & Generate" to enhance and generate in one step</p>
+            <p>4. Or check the enhancement checkbox and click Generate</p>
+        </div>
+        """
+    )
 
 if __name__ == "__main__":
     demo.launch(share=True)