Update app.py

app.py

@@ -4,45 +4,68 @@ import re
 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_prompt(text, original_user_input):
-    forbidden_words = ["4k", "8k", "high res", "high resolution", "hd", "ultra detailed", "masterpiece", "vray", "render"]
-    for word in forbidden_words:
-        text = re.sub(r'\b' + word + r'\b', "", text, flags=re.IGNORECASE)
+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"
+    ]
     
-    text = re.sub(r'\s+', ' ', text).strip()
+    # Clean the generated text
+    cleaned = generated_text
+    for word in bad_words:
+        # Remove the word (case insensitive)
+        cleaned = re.sub(r'\b' + word + r'\b', "", cleaned, flags=re.IGNORECASE)
     
-    if len(original_user_input.split()) < 4:
-        text = f"{original_user_input}, realistic look, centered in the image"
-    elif "centered" not in text.lower():
-        text += ", centered in the image"
-        
-    return text
-
-def generate_workflow(prompt, width, height, steps):
-    yield "🔍 Thinking (analysing AI)...", None, ""
+    # Remove extra commas and spaces created by removal
+    cleaned = re.sub(r',\s*,', ',', cleaned)
+    cleaned = re.sub(r'\s+', ' ', cleaned).strip().strip(',')
     
-    enhanced_results = enhancer_pipe(prompt, max_length=60, num_return_sequences=1)
-    raw_enhanced = enhanced_results[0]['generated_text']
+    # 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:
+        cleaned = f"{original_prompt}, detailed, centered in frame"
     
-    refined_prompt = clean_prompt(raw_enhanced, prompt)
+    return cleaned
+
+def generate_workflow(prompt, width, height, steps):
+    # Step 1: Analysis
+    yield "🔍 Thinking (Improving your prompt)...", None, ""
     
-    yield "🎨 Generating (Image generator AI)...", None, refined_prompt
+    # Generate extension
+    try:
+        # We limit max_length to keep it concise
+        enhanced_results = enhancer_pipe(prompt, max_length=60, 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
+        final_prompt = f"{prompt}, detailed, centered in frame"
+
+    # Step 2: Generation
+    yield "🎨 Generating (Drawing the image)...", None, final_prompt
     
     image = image_pipe(
-        prompt=refined_prompt,
+        prompt=final_prompt,
         width=int(width),
         height=int(height),
         num_inference_steps=int(steps),
@@ -51,16 +74,18 @@ def generate_workflow(prompt, width, height, steps):
         output_type="pil"
     ).images[0]
     
-    yield "✅ Ready", image, refined_prompt
+    # Step 3: Finish
+    yield "✅ Ready", image, final_prompt
 
-with gr.Blocks() as demo:
+# 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")
     
     with gr.Row():
         with gr.Column(scale=1):
             prompt_input = gr.Textbox(
                 label="💡 Your Idea", 
-                placeholder="e.g., Apple fruit",
+                placeholder="e.g., A cute dragon",
                 lines=3
             )
             
@@ -80,8 +105,7 @@ with gr.Blocks() as demo:
     generate_btn.click(
         fn=generate_workflow,
         inputs=[prompt_input, width_slider, height_slider, steps_slider],
-        outputs=[status_bar, image_output, refined_prompt_display],
-        api_name="predict"
+        outputs=[status_bar, image_output, refined_prompt_display]
     )
 
-demo.launch(theme=gr.themes.Soft(), title="AI Image Lab")
+demo.launch()