add default system prompt and refactor parameters for text generation

app.py

@@ -12,9 +12,28 @@ from PIL import Image
 hf_token = os.getenv("hf_token")
 login(token=hf_token)
 
+# Global constants and default values
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
-PRELOAD_MODELS = False  # Easy switch for preloading
+PRELOAD_MODELS = False
+
+# Default system prompt for text generation
+DEFAULT_SYSTEM_PROMPT = """Vous êtes un designer produit avec de solides connaissances dans la génération de texte en image. Vous recevrez une demande de produit sous forme de description succincte, et votre mission sera d'imaginer un nouveau design de produit répondant à ce besoin.
+
+Le livrable (réponse générée) sera exclusivement un texte de prompt pour l'IA de texte to image FLUX.1-schnell.
+
+Ce prompt devra inclure une description visuelle de l'objet mentionnant explicitement les aspects indispensables de sa fonction.
+A coté de ça vous devez aussi explicitement mentionner dans ce prompt les caractéristiques esthétiques/photo du rendu image (ex : photoréaliste, haute qualité, focale, grain, etc.), sachant que l'image sera l'image principale de cet objet dans le catalogue produit. Le fond de l'image générée doit être entièrement blanc.
+Le prompt doit être sans narration, peut être long mais ne doit pas dépasser 77 jetons."""
+
+# Default Flux parameters
+DEFAULT_SEED = 42
+DEFAULT_RANDOMIZE_SEED = True
+DEFAULT_WIDTH = 512
+DEFAULT_HEIGHT = 512
+DEFAULT_NUM_INFERENCE_STEPS = 6
+DEFAULT_GUIDANCE_SCALE = 0.0
+DEFAULT_TEMPERATURE = 0.9
 
 _text_gen_pipeline = None
 _image_gen_pipeline = None
@@ -64,15 +83,6 @@ def get_text_gen_pipeline():
             return None
     return _text_gen_pipeline
 
-# Default system prompt for text generation
-DEFAULT_SYSTEM_PROMPT = """Vous êtes un designer produit avec de solides connaissances dans la génération de texte en image. Vous recevrez une demande de produit sous forme de description succincte, et votre mission sera d'imaginer un nouveau design de produit répondant à ce besoin.
-
-Le livrable (réponse générée) sera exclusivement un texte de prompt pour l'IA de texte to image FLUX.1-schnell.
-
-Ce prompt devra inclure une description visuelle de l'objet mentionnant explicitement les aspects indispensables de sa fonction.
-A coté de ça vous devez aussi explicitement mentionner dans ce prompt les caractéristiques esthétiques/photo du rendu image (ex : photoréaliste, haute qualité, focale, grain, etc.), sachant que l'image sera l'image principale de cet objet dans le catalogue produit. Le fond de l'image générée doit être entièrement blanc.
-Le prompt doit être sans narration, peut être long mais ne doit pas dépasser 77 jetons."""
-
 @spaces.GPU()
 def refine_prompt(prompt, system_prompt=DEFAULT_SYSTEM_PROMPT, progress=gr.Progress()):
     text_gen = get_text_gen_pipeline()
@@ -114,12 +124,18 @@ def validate_dimensions(width, height):
     return True, None
 
 @spaces.GPU()
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)): 
+def infer(prompt, seed=DEFAULT_SEED, 
+          randomize_seed=DEFAULT_RANDOMIZE_SEED, 
+          width=DEFAULT_WIDTH, 
+          height=DEFAULT_HEIGHT, 
+          num_inference_steps=DEFAULT_NUM_INFERENCE_STEPS, 
+          progress=gr.Progress(track_tqdm=True)): 
     try:
         # Validate that prompt is not empty
         if not prompt or prompt.strip() == "":
             return None, "Please provide a valid prompt."
 
+        progress(0.1, desc="Loading model")
         pipe = get_image_gen_pipeline()
         if pipe is None:
             return None, "Image generation model is unavailable."
@@ -134,6 +150,7 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
         # Use default torch generator instead of cuda-specific generator
         generator = torch.Generator().manual_seed(seed)
         
+        progress(0.3, desc="Running inference")
         # Match the working example's parameters
         output = pipe(
             prompt=prompt, 
@@ -141,20 +158,23 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
             height=height,
             num_inference_steps=num_inference_steps, 
             generator=generator,
-            guidance_scale=0.0,  # Changed from 7.5 to 0.0
+            guidance_scale=DEFAULT_GUIDANCE_SCALE,
         )
             
+        progress(0.8, desc="Processing output")
         image = output.images[0]
+        progress(1.0, desc="Complete")
         return image, f"Image generated successfully with seed {seed}"
     except Exception as e:
         print(f"Error in infer: {str(e)}")  
         return None, f"Error generating image: {str(e)}"
  
-# Update examples to be a list of prompts only, not including other parameters
+ 
+# Format: [prompt, system_prompt]
 examples = [
-    "a backpack for kids, flower style",
-    "medieval flip flops",
-    "cat shaped cake mold",
+    ["a backpack for kids, flower style", DEFAULT_SYSTEM_PROMPT],
+    ["medieval flip flops", DEFAULT_SYSTEM_PROMPT],
+    ["cat shaped cake mold", DEFAULT_SYSTEM_PROMPT],
 ]
 
 css="""
@@ -165,26 +185,10 @@ css="""
 """
 
 def preload_models():
-    global _text_gen_pipeline, _image_gen_pipeline
-    
     print("Preloading models...")
-    success = True
-    
-    try:
-        _text_gen_pipeline = get_text_gen_pipeline()
-        if _text_gen_pipeline is None:
-            success = False
-    except Exception as e:
-        print(f"Error preloading text generation model: {str(e)}")
-        success = False
-    
-    try:
-        _image_gen_pipeline = get_image_gen_pipeline()
-        if _image_gen_pipeline is None:
-            success = False
-    except Exception as e:
-        print(f"Error preloading image generation model: {str(e)}")
-        success = False
+    text_success = get_text_gen_pipeline() is not None
+    image_success = get_image_gen_pipeline() is not None
+    success = text_success and image_success
     
     status = "Models preloaded successfully!" if success else "Error preloading models"
     print(status)
@@ -196,7 +200,6 @@ def preload_models():
 def process_example_pipeline(example_prompt, system_prompt=DEFAULT_SYSTEM_PROMPT, progress=gr.Progress()):
     # Step 1: Update status
     progress(0, desc="Starting example processing")
-    progress_status = "Selected example: " + example_prompt
     
     # Step 2: Refine the prompt
     progress(0.1, desc="Refining prompt with Mistral")
@@ -254,7 +257,7 @@ def create_interface():
                     # Mistral settings
                     temperature = gr.Slider(
                         label="Temperature",
-                        value=0.9,
+                        value=DEFAULT_TEMPERATURE,
                         minimum=0.0,
                         maximum=1.0,
                         step=0.05,
@@ -270,19 +273,19 @@ def create_interface():
                     
                 with gr.Tab("Flux"):
                     # Flux settings
-                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=DEFAULT_SEED)
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=DEFAULT_RANDOMIZE_SEED)
                     
                     with gr.Row():
-                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512)
-                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512)
+                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
+                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)
                     
                     num_inference_steps = gr.Slider(
                         label="Number of inference steps",
                         minimum=1,
                         maximum=50,
                         step=1,
-                        value=6,
+                        value=DEFAULT_NUM_INFERENCE_STEPS,
                     )
 
             # Examples section - simplified version that only updates the prompt fields