Update app.py

app.py

@@ -6,19 +6,14 @@ import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
-from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL, DPMSolverMultistepScheduler
+from diffusers import EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline
 from huggingface_hub import hf_hub_download, InferenceClient
+from diffusers import DiffusionPipeline
 
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-pipe = StableDiffusionXLPipeline.from_pretrained("SG161222/RealVisXL_V4.0", torch_dtype=torch.float16, vae=vae)
-pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
-pipe.to("cuda")
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
-refiner.to("cuda")
-
-pipe_fast = StableDiffusionXLPipeline.from_pretrained("SG161222/RealVisXL_V4.0_Lightning", torch_dtype=torch.float16, vae=vae, use_safetensors=True)
-pipe_fast.to("cuda")
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, revision="refs/pr/1").to(device)
 
 help_text = """
 To optimize image results:
@@ -51,11 +46,11 @@ pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_
 pipe_edit.to("cuda")
 
 client1 = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-system_instructions1 = "<|system|>\nAct as Image Prompt Generation expert, Your task is to modify prompt by USER to more better prompt for Image Generation in Stable Diffusion XL. \n Modify the user's prompt to generate a high-quality image by incorporating essential keywords and styles according to prompt if none style is mentioned than assume realistic. The optimized prompt may include keywords according to prompt for resolution (4K, HD, 16:9 aspect ratio, , etc.), image quality (cute, masterpiece, high-quality, vivid colors, intricate details, etc.), and desired art styles (realistic, anime, 3D, logo, futuristic, fantasy, etc.). Ensure the prompt is concise, yet comprehensive and choose keywords wisely, to generate an exceptional image that meets the user's expectations. \n Your task is to reply with final optimized prompt only. If you get big prompt make it concise. and Apply all keyword at last of prompt. Reply with optimized prompt only.\n<|user|>\n"
+system_instructions1 = "<|system|>\nAct as Image Prompt Generation expert, Your task is to modify prompt by USER to more better and detailed prompt for Image Generation. \n  Ensure the prompt is deatiled, yet descriptive to generate an exceptional image that meets the user's expectations. \n Your task is to reply with final optimized prompt only. Reply with optimized prompt only.\n<|user|>\n"
 
 def promptifier(prompt):
     formatted_prompt = f"{system_instructions1}{prompt}\n<|assistant|>\n"
-    stream = client1.text_generation(formatted_prompt, max_new_tokens=100)
+    stream = client1.text_generation(formatted_prompt, max_new_tokens=300)
     return stream
 
 # Generator
@@ -71,7 +66,6 @@ def king(type ,
         width: int = 1024,
         height: int = 1024,
         guidance_scale: float = 6,
-        fast=True,
         progress=gr.Progress(track_tqdm=True)
     ):
     if type=="Image Editing" :
@@ -102,22 +96,17 @@ def king(type ,
             print(f"BEFORE: {instruction} ")
             instruction = promptifier(instruction)
             print(f"AFTER: {instruction} ")
-        guidance_scale2=(guidance_scale/2)
-        if fast:
-            refine = pipe_fast(prompt = instruction,
-            guidance_scale = guidance_scale2, 
-            num_inference_steps = int(steps/2.5),
-            width = width, height = height,
-            generator = generator,
-            ).images[0]
-        else:            
-            image = pipe_fast( prompt = instruction,
-            negative_prompt=negative_prompt,
-            guidance_scale = guidance_scale, 
-            num_inference_steps = steps, 
-            width = width, height = height,
-            generator = generator, output_type="latent",
-            ).images 
+
+            image = pipe(
+                    prompt = instruction, 
+                    negative_prompt = negative_prompt,
+                    width = width,
+                    height = height,
+                    num_inference_steps = (steps/5), 
+                    generator = generator,
+                    guidance_scale=0.0,
+                    output_type="latent"
+            ).images
 
             refine = refiner( prompt=instruction,
                     negative_prompt = negative_prompt,
@@ -202,7 +191,6 @@ with gr.Blocks(css=css) as demo:
     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=False, scale=0)
-        fast = gr.Checkbox(label="FAST Generation", value=True, scale=0)
         
     with gr.Row():
         input_image = gr.Image(label="Image", type='filepath', interactive=True)
@@ -255,7 +243,6 @@ with gr.Blocks(css=css) as demo:
                 width,
                 height,
                 guidance_scale,
-                fast,
             ],
             outputs=[seed, input_image],
           api_name = "image_gen_pro",