Update app.py

app.py

@@ -188,7 +188,7 @@ def generate_caption(text_model, tokenizer, image_features, prompt_str: str, max
 @torch.no_grad()
 def stream_chat(input_image: Image.Image, caption_type: str, caption_tone: str, caption_length: str | int, lens_type: str = "", film_stock: str = "", composition_style: str = "", lighting_aspect: str = "", special_technique: str = "", color_effect: str = "") -> str:
     """
-    Generate a caption or style prompt based on the input image and parameters.
+    Generate a caption, training prompt, tags, or a style prompt for image generation based on the input image and parameters.
     """
     torch.cuda.empty_cache()
 
@@ -206,23 +206,39 @@ def stream_chat(input_image: Image.Image, caption_type: str, caption_tone: str,
     if prompt_key not in CAPTION_TYPE_MAP:
         raise ValueError(f"Invalid caption type: {prompt_key}")
 
-    prompt_str = CAPTION_TYPE_MAP[prompt_key][0].format(length=length, word_count=length)
-    
     if caption_type == "style_prompt":
-        # Extract the key from the dropdown value
-        lens_type_key = lens_type.split(":")[0].strip() if lens_type else ""
-        film_stock_key = film_stock.split(":")[0].strip() if film_stock else ""
-        composition_style_key = composition_style.split(":")[0].strip() if composition_style else ""
-        lighting_aspect_key = lighting_aspect.split(":")[0].strip() if lighting_aspect else ""
-        special_technique_key = special_technique.split(":")[0].strip() if special_technique else ""
-        color_effect_key = color_effect.split(":")[0].strip() if color_effect else ""
-
-        prompt_str += f" Lens type: {lens_type_key} ({lens_types_info[lens_type_key]}). " if lens_type_key else ""
-        prompt_str += f"Film stock: {film_stock_key} ({film_stocks_info[film_stock_key]}). " if film_stock_key else ""
-        prompt_str += f"Composition style: {composition_style_key} ({composition_styles_info[composition_style_key]}). " if composition_style_key else ""
-        prompt_str += f"Lighting aspect: {lighting_aspect_key} ({lighting_aspects_info[lighting_aspect_key]}). " if lighting_aspect_key else ""
-        prompt_str += f"Special technique: {special_technique_key} ({special_techniques_info[special_technique_key]}). " if special_technique_key else ""
-        prompt_str += f"Color effect: {color_effect_key} ({color_effects_info[color_effect_key]})." if color_effect_key else ""
+        # For style prompt, we'll create a custom prompt for the LLM
+        base_prompt = "Analyze the given image and create a detailed Stable Diffusion prompt for generating a new, creative image inspired by it. "
+        base_prompt += "The prompt should describe the main elements, style, and mood of the image, "
+        base_prompt += "but also introduce creative variations or enhancements. "
+        base_prompt += "Include specific details about the composition, lighting, and overall atmosphere. "
+        
+        # Add custom settings to the prompt
+        if lens_type:
+            lens_type_key = lens_type.split(":")[0].strip()
+            base_prompt += f"Incorporate the effect of a {lens_type_key} lens ({lens_types_info[lens_type_key]}). "
+        if film_stock:
+            film_stock_key = film_stock.split(":")[0].strip()
+            base_prompt += f"Apply the characteristics of {film_stock_key} film stock ({film_stocks_info[film_stock_key]}). "
+        if composition_style:
+            composition_style_key = composition_style.split(":")[0].strip()
+            base_prompt += f"Use a {composition_style_key} composition style ({composition_styles_info[composition_style_key]}). "
+        if lighting_aspect:
+            lighting_aspect_key = lighting_aspect.split(":")[0].strip()
+            base_prompt += f"Implement {lighting_aspect_key} lighting ({lighting_aspects_info[lighting_aspect_key]}). "
+        if special_technique:
+            special_technique_key = special_technique.split(":")[0].strip()
+            base_prompt += f"Apply the {special_technique_key} technique ({special_techniques_info[special_technique_key]}). "
+        if color_effect:
+            color_effect_key = color_effect.split(":")[0].strip()
+            base_prompt += f"Use a {color_effect_key} color effect ({color_effects_info[color_effect_key]}). "
+        
+        base_prompt += f"The final prompt should be approximately {length} words long. "
+        base_prompt += "Format the output as a single paragraph without numbering or bullet points."
+
+        prompt_str = base_prompt
+    else:
+        prompt_str = CAPTION_TYPE_MAP[prompt_key][0].format(length=length, word_count=length)
 
     # Debugging: Print the constructed prompt string
     print(f"Constructed Prompt: {prompt_str}")