Update app.py

app.py

@@ -6,7 +6,7 @@ import numpy as np
 
 import random
 import torch
-from diffusers import StableDiffusion3Pipeline, AutoencoderKL, StableDiffusionXLImg2ImgPipeline, EDMEulerScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler 
+from diffusers import StableDiffusion3Pipeline, AutoencoderKL, StableDiffusionXLImg2ImgPipeline, EulerAncestralDiscreteScheduler 
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
 from transformers import pipeline
@@ -55,14 +55,14 @@ 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("ford442/stable-diffusion-3.5-medium-bf16", torch_dtype=torch.bfloat16)
 #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++")
 
 #pipe.scheduler.config.requires_aesthetics_score = False
 #pipe.enable_model_cpu_offload()
-#pipe.to(device)
+pipe.to(device)
 #pipe = torch.compile(pipe)
 # pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear")
 
@@ -90,7 +90,7 @@ def filter_text(text):
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 4096
 
-@spaces.GPU(duration=80)
+@spaces.GPU(duration=90)
 def infer(
     prompt,
     negative_prompt,
@@ -139,7 +139,8 @@ def infer(
     print('-- filtered prompt --')
     print(enhanced_prompt)
     print('-- generating image --')
-    sd_image = pipe(
+    with torch.no_grad():
+      sd_image = pipe(
         prompt=enhanced_prompt,  # This conversion is fine
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
@@ -147,7 +148,7 @@ def infer(
         width=width,
         height=height,
         generator=generator
-    ).images[0]
+      ).images[0]
     print('-- got image --')
     image_path = f"sd35m_{seed}.png"
     sd_image.save(image_path,optimize=False,compress_level=0)