Update app.py

app.py

@@ -1,68 +1,59 @@
 import gradio as gr
 import torch
 import re
+import time
 from diffusers import DiffusionPipeline
 from transformers import pipeline, GPT2Tokenizer, GPT2LMHeadModel
 
-# 1. Setup Device
 device = "cpu"
 if torch.cuda.is_available():
     device = "cuda"
 
-# 2. Load Prompt Enhancer (The Brain)
 prompt_enhancer_id = "succinctly/text2image-prompt-generator"
 enhancer_tokenizer = GPT2Tokenizer.from_pretrained(prompt_enhancer_id)
 enhancer_model = GPT2LMHeadModel.from_pretrained(prompt_enhancer_id)
 enhancer_pipe = pipeline("text-generation", model=enhancer_model, tokenizer=enhancer_tokenizer, device=device)
 
-# 3. Load Image Generator (The Artist) - CPU Optimized
 image_model_id = "SimianLuo/LCM_Dreamshaper_v7"
 image_pipe = DiffusionPipeline.from_pretrained(image_model_id)
 image_pipe.to(device)
 
 def clean_and_format_prompt(generated_text, original_prompt):
-    # List of "filler" words to remove
     bad_words = [
         "4k", "8k", "high definition", "high res", "high resolution", 
         "hd", "ultra detailed", "masterpiece", "photorealistic", 
         "best quality", "vray", "unreal engine", "octane render"
     ]
     
-    # Clean the generated text
     cleaned = generated_text
+    instruction_trigger = "Enhanced prompt:"
+    if instruction_trigger in cleaned:
+        cleaned = cleaned.split(instruction_trigger)[-1]
+
     for word in bad_words:
-        # Remove the word (case insensitive)
         cleaned = re.sub(r'\b' + word + r'\b', "", cleaned, flags=re.IGNORECASE)
     
-    # Remove extra commas and spaces created by removal
     cleaned = re.sub(r',\s*,', ',', cleaned)
     cleaned = re.sub(r'\s+', ' ', cleaned).strip().strip(',')
     
-    # Logic: If the enhancer didn't add much substance, use a template
-    # This prevents "Apple" -> "Apple 4k 8k" (which becomes just "Apple" after cleaning)
-    if len(cleaned) < len(original_prompt) + 10:
+    if len(cleaned) < 5:
         cleaned = f"{original_prompt}, detailed, centered in frame"
     
     return cleaned
 
 def generate_workflow(prompt, width, height, steps):
-    # Step 1: Analysis
-    yield "🔍 Thinking (Improving your prompt)...", None, ""
+    start_time = time.time()
+    yield "🔍 Thinking (analysing AI)...", None, ""
     
-    # Generate extension
     try:
-        # We limit max_length to keep it concise
-        enhanced_results = enhancer_pipe(prompt, max_length=60, num_return_sequences=1)
+        instructional_prompt = f"Enhance the user prompt so it is suitable for an image generator, and focus on the object, not on the quality, resolution etc. User prompt: {prompt}. Enhanced prompt:"
+        enhanced_results = enhancer_pipe(instructional_prompt, max_new_tokens=40, num_return_sequences=1)
         raw_output = enhanced_results[0]['generated_text']
-        
-        # Apply our cleaning logic
         final_prompt = clean_and_format_prompt(raw_output, prompt)
-    except Exception as e:
-        # Fallback if enhancer fails
+    except:
         final_prompt = f"{prompt}, detailed, centered in frame"
 
-    # Step 2: Generation
-    yield "🎨 Generating (Drawing the image)...", None, final_prompt
+    yield "🎨 Generating (Image generator AI)...", None, final_prompt
     
     image = image_pipe(
         prompt=final_prompt,
@@ -74,10 +65,11 @@ def generate_workflow(prompt, width, height, steps):
         output_type="pil"
     ).images[0]
     
-    # Step 3: Finish
-    yield "✅ Ready", image, final_prompt
+    end_time = time.time()
+    duration = round(end_time - start_time, 2)
+    
+    yield f"✅ Done in {duration}s", image, final_prompt
 
-# UI Setup - Title and Theme belong here!
 with gr.Blocks(theme=gr.themes.Soft(), title="AI Image Lab") as demo:
     gr.Markdown("# 🎨 AI Image Lab")
     
@@ -93,7 +85,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI Image Lab") as demo:
                 width_slider = gr.Slider(256, 768, 512, step=64, label="📏 Width")
                 height_slider = gr.Slider(256, 768, 512, step=64, label="📐 Height")
             
-            steps_slider = gr.Slider(1, 15, 4, step=1, label="🏃 Steps")
+            steps_slider = gr.Slider(4, 12, 5, step=1, label="🏃 Steps")
             
             generate_btn = gr.Button("🚀 Generate", variant="primary")