Update app.py

app.py

@@ -17,11 +17,18 @@ 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"]
+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(word, "", text, flags=re.IGNORECASE)
+        text = re.sub(r'\b' + word + r'\b', "", text, flags=re.IGNORECASE)
+    
     text = re.sub(r'\s+', ' ', text).strip()
+    
+    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):
@@ -30,11 +37,8 @@ def generate_workflow(prompt, width, height, steps):
     enhanced_results = enhancer_pipe(prompt, max_length=60, num_return_sequences=1)
     raw_enhanced = enhanced_results[0]['generated_text']
     
-    refined_prompt = clean_prompt(raw_enhanced)
+    refined_prompt = clean_prompt(raw_enhanced, 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(
@@ -49,7 +53,7 @@ def generate_workflow(prompt, width, height, steps):
     
     yield "✅ Ready", image, refined_prompt
 
-with gr.Blocks(theme=gr.themes.Soft(), title="AI Image Lab") as demo:
+with gr.Blocks() as demo:
     gr.Markdown("# 🎨 AI Image Lab")
     
     with gr.Row():
@@ -80,4 +84,4 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI Image Lab") as demo:
         api_name="predict"
     )
 
-demo.launch(show_api=True)
+demo.launch(theme=gr.themes.Soft(), title="AI Image Lab")