Update app.py

app.py

@@ -17,7 +17,7 @@ from typing import Tuple
 import paramiko
 import datetime
 #import diffusers
-from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL, EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, UNet2DConditionModel, AutoencoderKL, EulerAncestralDiscreteScheduler
 from diffusers.models.attention_processor import Attention, AttnProcessor2_0
 from transformers import CLIPTextModelWithProjection, CLIPTextModel, CLIPTokenizer
 
@@ -105,6 +105,7 @@ def load_and_prepare_model():
         unet=unet,
         vae=None,
     )
+    
     #pipe.scheduler=scheduler
     #pipe.tokenizer=tokenizer_1
     #pipe.tokenizer_2=tokenizer_2
@@ -117,13 +118,24 @@ def load_and_prepare_model():
     pipe.watermark=None
     pipe.safety_checker=None
     #pipe.unet.to(memory_format=torch.channels_last)
-    pipe.to(device=device, dtype=torch.bfloat16)
+    pipe.to(device=device)
     pipe.vae = vaeXL.to(device) #.to('cpu') #.to(torch.bfloat16)
     pipe.vae.set_default_attn_processor()
     return pipe
     
 pipe = load_and_prepare_model()
 
+sched = EulerAncestralDiscreteScheduler.from_pretrained('ford442/stable-diffusion-xl-refiner-1.0-bf16', subfolder='scheduler',beta_schedule="scaled_linear", token=HF_TOKEN)
+
+refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+        "ford442/stable-diffusion-xl-refiner-1.0-bf16",
+        requires_aesthetics_score=True,
+        token=HF_TOKEN
+)
+
+refiner.vae=vaeX
+refiner.to(device=device)
+
 MAX_SEED = np.iinfo(np.int32).max
 
 neg_prompt_2 = " 'non-photorealistic':1.5, 'unrealistic skin','unattractive face':1.3, 'low quality':1.1, ('dull color scheme', 'dull colors', 'digital noise':1.2),'amateurish', 'poorly drawn face':1.3, 'poorly drawn', 'distorted face', 'low resolution', 'simplistic' "
@@ -196,15 +208,31 @@ def generate_30(
         "guidance_scale": guidance_scale,
         "num_inference_steps": num_inference_steps,
         "generator": generator,
-        "output_type": "pil",
+        "output_type": "latent",
+        "denoising_end": 0.75,
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
     images = []
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
     batch_options = options.copy()
-    rv_image = pipe(**batch_options).images[0]
+    rv_image = pipe(**batch_options).images
+    options = {
+        "prompt": [prompt],
+        "image": rv_image,
+        "negative_prompt": [negative_prompt],
+        "negative_prompt_2": [neg_prompt_2],
+        "width": width,
+        "height": height,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "generator": generator,
+        "output_type": "latent",
+        "denoising_start": 0.75,
+    } 
+    batch_options = options.copy()
+    rv_image = refiner(**batch_options).images[0]
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     sd_image_path = f"rv50_G_{timestamp}.png"
     rv_image.save(sd_image_path,optimize=False,compress_level=0)
     upload_to_ftp(sd_image_path)    
@@ -238,15 +266,31 @@ def generate_60(
         "guidance_scale": guidance_scale,
         "num_inference_steps": num_inference_steps,
         "generator": generator,
-        "output_type": "pil",
+        "output_type": "latent",
+        "denoising_end": 0.75,
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
     images = []
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
     batch_options = options.copy()
-    rv_image = pipe(**batch_options).images[0]
+    rv_image = pipe(**batch_options).images
+    options = {
+        "prompt": [prompt],
+        "image": rv_image,
+        "negative_prompt": [negative_prompt],
+        "negative_prompt_2": [neg_prompt_2],
+        "width": width,
+        "height": height,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "generator": generator,
+        "output_type": "latent",
+        "denoising_start": 0.75,
+    } 
+    batch_options = options.copy()
+    rv_image = refiner(**batch_options).images[0]
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     sd_image_path = f"rv50_G_{timestamp}.png"
     rv_image.save(sd_image_path,optimize=False,compress_level=0)
     upload_to_ftp(sd_image_path)    
@@ -280,15 +324,31 @@ def generate_90(
         "guidance_scale": guidance_scale,
         "num_inference_steps": num_inference_steps,
         "generator": generator,
-        "output_type": "pil",
+        "output_type": "latent",
+        "denoising_end": 0.75,
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
     images = []
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
     batch_options = options.copy()
-    rv_image = pipe(**batch_options).images[0]
+    rv_image = pipe(**batch_options).images
+    options = {
+        "prompt": [prompt],
+        "image": rv_image,
+        "negative_prompt": [negative_prompt],
+        "negative_prompt_2": [neg_prompt_2],
+        "width": width,
+        "height": height,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "generator": generator,
+        "output_type": "latent",
+        "denoising_start": 0.75,
+    } 
+    batch_options = options.copy()
+    rv_image = refiner(**batch_options).images[0]
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     sd_image_path = f"rv50_G_{timestamp}.png"
     rv_image.save(sd_image_path,optimize=False,compress_level=0)
     upload_to_ftp(sd_image_path)