Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+import re
 from diffusers import DiffusionPipeline
 from transformers import pipeline, GPT2Tokenizer, GPT2LMHeadModel
 
@@ -16,13 +17,25 @@ image_model_id = "SimianLuo/LCM_Dreamshaper_v7"
 image_pipe = DiffusionPipeline.from_pretrained(image_model_id)
 image_pipe.to(device)
 
+def clean_prompt(text):
+    forbidden_words = ["4k", "8k", "high res", "high resolution", "hd", "ultra detailed", "masterpiece"]
+    for word in forbidden_words:
+        text = re.sub(word, "", text, flags=re.IGNORECASE)
+    text = re.sub(r'\s+', ' ', text).strip()
+    return text
+
 def generate_workflow(prompt, width, height, steps):
-    yield "Thinking (analysing AI)...", None, ""
+    yield "🔍 Thinking (analysing AI)...", None, ""
+    
+    enhanced_results = enhancer_pipe(prompt, max_length=60, num_return_sequences=1)
+    raw_enhanced = enhanced_results[0]['generated_text']
     
-    enhanced_results = enhancer_pipe(prompt, max_length=75, num_return_sequences=1)
-    refined_prompt = enhanced_results[0]['generated_text']
+    refined_prompt = clean_prompt(raw_enhanced)
     
-    yield "Generating (Image generator AI)...", None, refined_prompt
+    if "centered" not in refined_prompt.lower():
+        refined_prompt += ", centered in the image"
+        
+    yield "🎨 Generating (Image generator AI)...", None, refined_prompt
     
     image = image_pipe(
         prompt=refined_prompt,
@@ -34,36 +47,37 @@ def generate_workflow(prompt, width, height, steps):
         output_type="pil"
     ).images[0]
     
-    yield "Ready", image, refined_prompt
+    yield "✅ Ready", image, refined_prompt
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# AI Image Lab")
+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., a lonely robot on mars",
+                label="💡 Your Idea", 
+                placeholder="e.g., Apple fruit",
                 lines=3
             )
             
             with gr.Row():
-                width_slider = gr.Slider(256, 768, 512, step=64, label="Width")
-                height_slider = gr.Slider(256, 768, 512, step=64, label="Height")
+                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="Inference Steps")
+            steps_slider = gr.Slider(1, 15, 4, step=1, label="🏃 Steps")
             
-            generate_btn = gr.Button("Generate", variant="primary")
+            generate_btn = gr.Button("🚀 Generate", variant="primary")
             
         with gr.Column(scale=1):
             status_bar = gr.Markdown("### Status: **Ready**")
-            image_output = gr.Image(label="Result")
-            refined_prompt_display = gr.Textbox(label="Enhanced Prompt Used", interactive=False)
+            image_output = gr.Image(label="🖼️ Result")
+            refined_prompt_display = gr.Textbox(label="📝 Enhanced Prompt Used", interactive=False)
 
     generate_btn.click(
         fn=generate_workflow,
         inputs=[prompt_input, width_slider, height_slider, steps_slider],
-        outputs=[status_bar, image_output, refined_prompt_display]
+        outputs=[status_bar, image_output, refined_prompt_display],
+        api_name="predict"
     )
 
-demo.launch()
+demo.launch(show_api=True)