Update app.py

app.py

@@ -51,8 +51,8 @@ pipe_edit.to("cuda")
 
 def promptifier(prompt):
     client1 = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
-    system_instructions1 = "<s>[SYSTEM] Your task is to modify prompt by USER to more better prompt for Image Generation in Stable Diffusion XL, reply with prompt only, Your task is to reply with final prompt in SDXL image generation format only. Just reply with pure prompt.[USER]"
-    formatted_prompt = f"{system_instructions1} {prompt} [FINAL_PROMPT]"
+    system_instructions1 = "<s>[SYSTEM] Your task is to modify prompt by USER to more better prompt for Image Generation in Stable Diffusion XL, you have to optiomize prompt and also add some keywords like, 4k, realistic, featuristic according to prompt and also break prompt into sub-lines using comma, Your task is to reply with final optimized prompt only. Just reply with prompt only.[USER]"
+    formatted_prompt = f"{system_instructions1} {prompt} [PROMPT]"
     stream = client1.text_generation(formatted_prompt, max_new_tokens=80, stream=True, details=True, return_full_text=False)
     return "".join([response.token.text for response in stream if response.token.text != "</s>"])
 
@@ -95,15 +95,19 @@ def king(type ,
             seed = random.randint(0, 999999)
         generator = torch.Generator().manual_seed(seed)
         if fast:
+            if enhance_prompt:
+                print(f"BEFORE: {instruction} ")
+                instruction = promptifier(instruction)
+                print(f"AFTER: {instruction} ")
             steps=int(steps/2.5)
-            guidance_scale2=(guidance_scale/3)
+            guidance_scale2=(guidance_scale/2)
             
             refine = pipe_fast( prompt = instruction,
             guidance_scale = guidance_scale2, 
             num_inference_steps = steps, 
             width = width, height = height,
-            generator = generator,
-            ).images[0]
+            generator = generator, output_type="latent",
+            ).images
         else:
             if enhance_prompt:
                 print(f"BEFORE: {instruction} ")
@@ -119,12 +123,12 @@ def king(type ,
             generator = generator, output_type="latent",
             ).images 
 
-            refine = refiner( prompt=instruction,
+        refine = refiner( prompt=instruction,
                 negative_prompt = negative_prompt,
                 guidance_scale = guidance_scale,
                 num_inference_steps=  steps,
                 image=image, generator=generator,
-            ).images[0]        
+        ).images[0]        
         return seed, refine
 
 client = InferenceClient()