Update app.py

app.py

@@ -84,8 +84,8 @@ def load_text_summarization_model():
 tokenizer, model = load_text_summarization_model()
 
 
-@spaces.GPU()
-def generate_image_with_flux(
+#@spaces.GPU()
+def generate_image_with_flux_old(
     text: str,
     seed: int = 42,
     width: int = 1024,
@@ -109,7 +109,6 @@ def generate_image_with_flux(
     # Initialize FLUX pipeline here
     dtype = torch.bfloat16
     device = "cuda" if torch.cuda.is_available() else "cpu"
-
     torch.cuda.empty_cache()  # Clear cache
     gc.collect()  # Run garbage collection
     flux_pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
@@ -128,6 +127,37 @@ def generate_image_with_flux(
     print("DEBUG: Image generated successfully.")
     return image
 
+@spaces.GPU()
+def generate_image_with_flux(
+    text: str,
+    seed: int = 42,
+    width: int = 1024,
+    height: int = 1024,
+    num_inference_steps: int = 4,
+    randomize_seed: bool = True):
+    """
+    Generates an image from text using FLUX.
+    """
+    print(f"DEBUG: Generating image with FLUX for text: '{text}'")
+    
+    # Use the global flux_pipe (which was already initialized at startup)
+    global flux_pipe
+    if flux_pipe is None:
+        raise RuntimeError("FLUX pipeline not initialized because CUDA is unavailable.")
+    
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)  # Specify device for generator
+    image = flux_pipe(
+        prompt=text,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        guidance_scale=0.0
+    ).images[0]
+    print("DEBUG: Image generated successfully.")
+    return image
 # --------- End of MinDalle Functions ---------
 # Merge audio files