Update app.py

app.py

@@ -22,8 +22,10 @@ if not torch.cuda.is_available():
 
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "3000"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
+ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
+PORT = int(os.getenv("DEMO_PORT", "15432"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -105,6 +107,12 @@ if torch.cuda.is_available():
         print("Using DALL-E 3 Consistency Decoder")
         pipe.vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=torch.float16)
 
+    if ENABLE_CPU_OFFLOAD:
+        pipe.enable_model_cpu_offload()
+    else:
+        pipe.to(device)
+        print("Loaded on Device!")
+        
     # speed-up T5
     pipe.text_encoder.to_bettertransformer()
 
@@ -125,6 +133,9 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     return seed
 
 
+@torch.no_grad()
+@torch.inference_mode()
+@spaces.GPU(duration=30)
 def generate(
         prompt: str,
         negative_prompt: str = "",