Update app.py

app.py

@@ -10,7 +10,7 @@ from diffusers import DiffusionPipeline, StableDiffusionXLPipeline, EDMEulerSche
 from huggingface_hub import hf_hub_download, InferenceClient
 
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-pipe = StableDiffusionXLPipeline.from_pretrained("fluently/Fluently-XL-Final", torch_dtype=torch.float16, vae=vae)
+pipe = StableDiffusionXLPipeline.from_pretrained("SG161222/RealVisXL_V3.0", torch_dtype=torch.float16, vae=vae)
 pipe.load_lora_weights("KingNish/Better-Image-XL-Lora", weight_name="example-03.safetensors", adapter_name="lora")
 pipe.set_adapters("lora")
 pipe.to("cuda")
@@ -51,19 +51,28 @@ pipe_edit = StableDiffusionXLInstructPix2PixPipeline.from_single_file( edit_file
 pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
 pipe_edit.to("cuda")
 
+def promptifier(prompt):
+    client1 = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
+    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. [USER]"
+    formatted_prompt = f"{system_instructions1} {prompt} [FINAL_PROMPT]"
+    stream = client1.text_generation(formatted_prompt, max_new_tokens=80, stream=True, details=True, return_full_text=False)
+    instructions = "".join([response.token.text for response in stream if response.token.text != "</s>"])
+    return instructions
+
 # Generator
 @spaces.GPU(duration=30, queue=False)
 def king(type ,
         input_image ,
         instruction: str ,
         negative_prompt: str ="",
+        enhance_prompt: bool = True
         steps: int = 25,
         randomize_seed: bool = False,
         seed: int = 2404,
         width: int = 1024,
         height: int = 1024,
         guidance_scale: float = 6,
-        fast=True,
+        fast=False,
         use_resolution_binning: bool = True,
         progress=gr.Progress(track_tqdm=True),
     ):
@@ -90,18 +99,27 @@ def king(type ,
             seed = random.randint(0, 999999)
         generator = torch.Generator().manual_seed(seed)
         if fast:
-            pipes=pipe_fast
             steps=int(steps/2.5)
-            guidance_scale=(guidance_scale/3)
+            guidance_scale2=(guidance_scale/3)
+            
+            image = pipe_fast( prompt = instruction,
+            guidance_scale = guidance_scale2, 
+            num_inference_steps = steps, 
+            width = width, height = height,
+            generator = generator, output_type="latent",
+            ).images 
         else:
-            pipes=pipe
-        image = pipes( prompt = instruction,
+            guidance_scale2=(guidance_scale/2) 
+            if enhance_prompt:
+                instruction = promptifier(instruction)
+                
+            image = pipe( prompt = instruction,
             negative_prompt=negative_prompt,
-            guidance_scale = guidance_scale, 
+            guidance_scale = guidance_scale2, 
             num_inference_steps = steps, 
             width = width, height = height,
             generator = generator, output_type="latent",
-        ).images         
+            ).images 
 
         refine = refiner( prompt=instruction,
             negative_prompt = negative_prompt,
@@ -180,9 +198,9 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         generate_button = gr.Button("Run", scale=0)
     with gr.Row():
         type = gr.Dropdown(["Image Generation","Image Editing"], label="Task", value="Image Generation",interactive=True)
-        enhance_prompt = gr.Checkbox(label="Enhance prompt", value = True, scale=0)
-        fast = gr.Checkbox(label="FAST Gen", value=True, scale=0)
-
+        enhance_prompt = gr.Checkbox(label="Enhance prompt", value=True, scale=0)
+        fast = gr.Checkbox(label="FAST Generation", value=False, scale=0)
+        
     with gr.Row():
         input_image = gr.Image(label="Image", type='filepath', interactive=True)
 
@@ -235,6 +253,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 input_image,
                 instruction,
                 negative_prompt,
+                enhance_prompt,
                 steps,
                 randomize_seed,
                 seed,