Change response parsing

app.py

@@ -52,7 +52,7 @@ def create_prompt_for_image_generation(user_prompt: str) -> str:
 	3.	Color Palette: Include details about the dominant colors or overall color scheme.
 	4.	Perspective and Camera Details: Mention the point of view (e.g., wide-angle, close-up), camera type, lens, aperture, and lighting conditions if applicable.
 	5.	Additional Details: Highlight any specific objects, text, or unique features with clear emphasis (e.g., 'with green text' or 'emphasis on golden hour lighting').
-	6.	Output Settings: Suggest aspect ratio, output format (e.g., PNG), quality level, and seed for reproducibility.
+	6.	Output Settings: Suggest aspect ratio, output format (e.g., JPG), quality level, and seed for reproducibility.
     Ensure that the generated prompt is logical, descriptive, and written in natural language to maximize compatibility with FLUX-Schnell capabilities.
     Example Input:
     An image of a serene forest with a small cabin.
@@ -64,7 +64,7 @@ def create_prompt_for_image_generation(user_prompt: str) -> str:
     The color palette includes rich greens, warm browns for the cabin, and soft gray mist. 
     The perspective is slightly elevated as if viewed from a drone camera at sunrise, 
     capturing golden hour lighting for soft shadows and warm highlights. 
-    The aspect ratio is 1:1, using seed 42 for reproducibility.
+    The aspect ratio is 1:1, output format JPG, high quality, using seed 42 for reproducibility.
     """
     model = HfApiModel(
     max_tokens=384,
@@ -80,8 +80,8 @@ def create_prompt_for_image_generation(user_prompt: str) -> str:
         # response = model(
         #     prompt=prompt, temperature=1., max_tokens=512)
         response = model(messages, stop_sequences=["END"])
-        # return response['choices'][0]['text'] 
-        return response
+        return response['choices'][0]['text'] 
+        # return response
 
     except Exception as e:
         print(f"Error during LLM call: {str(e)}")