Update app.py

app.py

@@ -55,8 +55,8 @@ checkpoint = "microsoft/Phi-3.5-mini-instruct"
 #vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", torch_dtype=torch.bfloat16, device_map='balanced')
 
-#pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16", torch_dtype=torch.bfloat16, device_map='balanced')
-pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", token=hftoken, torch_dtype=torch.float32, device_map='balanced')
+pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16", torch_dtype=torch.bfloat16, device_map='balanced')
+#pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", token=hftoken, torch_dtype=torch.float32, device_map='balanced')
 
 # pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
 
@@ -90,7 +90,7 @@ def filter_text(text):
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 4096
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=80)
 def infer(
     prompt,
     negative_prompt,
@@ -103,7 +103,7 @@ def infer(
     progress=gr.Progress(track_tqdm=True),
 ):
     seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device='cpu').manual_seed(seed)
+    generator = torch.Generator(device='cuda').manual_seed(seed)
 
     system_prompt_rewrite = (
     "You are an AI assistant that rewrites image prompts to be more descriptive and detailed."
@@ -150,7 +150,7 @@ def infer(
     ).images[0]
     print('-- got image --')
     image_path = f"sd35m_{seed}.png"
-    sd_image.save(image_path)
+    sd_image.save(image_path,optimize=False,compress_level=0))
     upload_to_ftp(image_path)
     refine = refiner(
             prompt=f"{prompt}, high quality masterpiece, complex details",
@@ -161,7 +161,7 @@ def infer(
             generator=generator,
     ).images[0]  
     refine_path = f"refine_{seed}.png"
-    refine.save(refine_path)
+    refine.save(refine_path,optimize=False,compress_level=0))
     upload_to_ftp(refine_path)
     return refine, seed, refine_path, enhanced_prompt