Update app.py

app.py

@@ -33,7 +33,7 @@ import gradio as gr
 from lxml.html import fromstring
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file, save_file
-from diffusers import FluxPipeline
+from diffusers import DiffusionPipeline
 from PIL import Image, ImageDraw, ImageFont
 from transformers import pipeline, T5ForConditionalGeneration, T5Tokenizer
 from refiners.fluxion.utils import manual_seed
@@ -418,7 +418,7 @@ CHECKPOINTS = ESRGANUpscalerCheckpoints(
 )
 
 device = DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
+DTYPE = dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
 enhancer = ESRGANUpscaler(checkpoints=CHECKPOINTS, device=DEVICE, dtype=DTYPE)
 
 # logging
@@ -434,15 +434,15 @@ root.addHandler(handler)
 
 # constant data
 
-base = "black-forest-labs/FLUX.1-schnell"
+MAX_SEED = np.iinfo(np.int32).max
 
 # precision data
 
-seq=256
+seq=512
 width=1536
 height=1536
-image_steps=8
-img_accu=0
+image_steps=50
+img_accu=9.0
 
 # ui data
 
@@ -502,7 +502,9 @@ function custom(){
 
 # torch pipes
 
-image_pipe = FluxPipeline.from_pretrained(base, torch_dtype=torch.bfloat16).to(device)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
+good_vae = AutoencoderKL.from_pretrained("ostris/Flex.1-alpha", subfolder="vae", torch_dtype=dtype).to(device)
+image_pipe = DiffusionPipeline.from_pretrained("ostris/Flex.1-alpha", torch_dtype=dtype, vae=taef1).to(device)
 image_pipe.enable_model_cpu_offload()
 image_pipe.enable_vae_slicing()
 image_pipe.enable_vae_tiling()
@@ -511,7 +513,7 @@ image_pipe.enable_vae_tiling()
 
 def upscaler(
     input_image: Image.Image,
-    prompt: str = "Photorealistic, Hyperrealistic, Realistic Photography, High-Quality Photography, Natural.",
+    prompt: str = "Hyper realistic photography, Natural visual content.",
     negative_prompt: str = "Distorted, Discontinuous, Blurry, Doll-Like, Overly-Plastic, Low-Quality, Painted, Smoothed, Artificial, Phony, Gaudy, Digital Effects.",
     seed: int = int(str(random.random()).split(".")[1]),
     upscale_factor: int = 2,
@@ -613,13 +615,16 @@ def pipe_generate_image(p1,p2):
     imgs = image_pipe(
             prompt=p1,
             negative_prompt=p2,
+            progress=gr.Progress(track_tqdm=True),
             height=height,
             width=width,
+            safety_checker=None,
             guidance_scale=img_accu,
             num_images_per_prompt=1,
             num_inference_steps=image_steps,
             max_sequence_length=seq,
-            generator=torch.Generator(device).manual_seed(int(str(random.random()).split(".")[1]))
+            good_vae=good_vae,
+            generator=torch.Generator(device).manual_seed(random.randint(0, MAX_SEED))
     ).images
     log(f'RET pipe_generate')
     return imgs