Update app.py

app.py

@@ -3,57 +3,36 @@ import torch
 import re
 import time
 from diffusers import DiffusionPipeline
-from transformers import pipeline, GPT2Tokenizer, GPT2LMHeadModel
+# from transformers import pipeline, GPT2Tokenizer, GPT2LMHeadModel
 
 device = "cpu"
 if torch.cuda.is_available():
     device = "cuda"
 
-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)
+# 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)
 
 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):
-    bad_words = [
-        "4k", "8k", "high definition", "high res", "high resolution", 
-        "hd", "ultra detailed", "masterpiece", "photorealistic", 
-        "best quality", "vray", "unreal engine", "octane render"
-    ]
-    
-    cleaned = generated_text
-    instruction_trigger = "Enhanced prompt:"
-    if instruction_trigger in cleaned:
-        cleaned = cleaned.split(instruction_trigger)[-1]
-
-    for word in bad_words:
-        cleaned = re.sub(r'\b' + word + r'\b', "", cleaned, flags=re.IGNORECASE)
-    
-    cleaned = re.sub(r',\s*,', ',', cleaned)
-    cleaned = re.sub(r'\s+', ' ', cleaned).strip().strip(',')
-    
-    if len(cleaned) < 5:
-        cleaned = f"{original_prompt}, detailed, centered in frame"
-    
-    return cleaned
-
 def generate_workflow(prompt, width, height, steps):
     start_time = time.time()
-    yield "🔍 Thinking (analysing AI)...", None, ""
     
-    try:
-        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']
-        final_prompt = clean_and_format_prompt(raw_output, prompt)
-    except:
-        final_prompt = f"{prompt}, detailed, centered in frame"
+    # yield "🔍 Thinking (analysing AI)...", None
+    
+    # try:
+    #     instructional_prompt = f"Enhance the user prompt... {prompt}"
+    #     enhanced_results = enhancer_pipe(instructional_prompt, max_new_tokens=40, num_return_sequences=1)
+    #     final_prompt = clean_and_format_prompt(enhanced_results[0]['generated_text'], prompt)
+    # except:
+    #     final_prompt = f"{prompt}, centered and realistic (if applicable)"
+
+    final_prompt = f"{prompt}, centered and realistic (if applicable)"
 
-    yield "🎨 Generating (Image generator AI)...", None, final_prompt
+    yield "🎨 Generating (Image generator AI)...", None
     
     image = image_pipe(
         prompt=final_prompt,
@@ -68,7 +47,7 @@ def generate_workflow(prompt, width, height, steps):
     end_time = time.time()
     duration = round(end_time - start_time, 2)
     
-    yield f"✅ Done in {duration}s", image, final_prompt
+    yield f"✅ Done in {duration}s", image
 
 with gr.Blocks(theme=gr.themes.Soft(), title="AI Image Lab") as demo:
     gr.Markdown("# 🎨 AI Image Lab")
@@ -92,12 +71,12 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI Image Lab") as demo:
         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)
+            # 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]
     )
 
 demo.launch()