Update app.py

app.py

@@ -21,7 +21,7 @@ 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", "6000"))
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
@@ -105,7 +105,7 @@ 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:
+if ENABLE_CPU_OFFLOAD:
         pipe.enable_model_cpu_offload()
     else:
         pipe.to(device)
@@ -118,33 +118,35 @@ if torch.cuda.is_available():
         pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=True)
         print("Model Compiled!")
 
+
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
     return unique_name
 
+
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
 
+
 def generate(
         prompt: str,
         negative_prompt: str = "",
         style: str = DEFAULT_STYLE_NAME,
         use_negative_prompt: bool = False,
-        num_imgs: int = 1,
         seed: int = 0,
         width: int = 1024,
         height: int = 1024,
-        num_inference_steps: int = 4,
+        inference_steps: int = 4,
         randomize_seed: bool = False,
         use_resolution_binning: bool = True,
         progress=gr.Progress(track_tqdm=True),
 ):
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator().manual_seed(seed)
-
+    
     if not use_negative_prompt:
         negative_prompt = None  # type: ignore
     prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
@@ -290,6 +292,7 @@ with gr.Blocks() as demo:
             negative_prompt,
             style_selection,
             use_negative_prompt,
+            num_imgs,
             seed,
             width,
             height,
@@ -300,6 +303,6 @@ with gr.Blocks() as demo:
         api_name="run",
     )
 
-
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()
+    # demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=11900, debug=True)