Update app.py

app.py

@@ -40,9 +40,8 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
 
 from diffusers import DiffusionPipeline
 
-# Use torch.float16 for better stability on NVIDIA GPUs
 base_model = "black-forest-labs/FLUX.1-dev"
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
 lora_repo = "strangerzonehf/Flux-Super-Realism-LoRA"
 trigger_word = "Super Realism"  # Leave blank if no trigger word is needed.
 pipe.load_lora_weights(lora_repo)
@@ -90,17 +89,23 @@ def generate_image_flux(
     positive_prompt = apply_style(style_name, prompt)
     if trigger_word:
         positive_prompt = f"{trigger_word} {positive_prompt}"
-    # Clear GPU cache before generating to help avoid memory fragmentation errors.
+    
+    # Clear cache before generation
     torch.cuda.empty_cache()
-    images = pipe(
-        prompt=positive_prompt,
-        width=width,
-        height=height,
-        guidance_scale=guidance_scale,
-        num_inference_steps=28,
-        num_images_per_prompt=1,
-        output_type="pil",
-    ).images
+    
+    # Wrap the pipeline call in no_grad and autocast contexts
+    with torch.no_grad():
+        with torch.cuda.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
+            outputs = pipe(
+                prompt=positive_prompt,
+                width=width,
+                height=height,
+                guidance_scale=guidance_scale,
+                num_inference_steps=28,
+                num_images_per_prompt=1,
+                output_type="pil",
+            )
+    images = outputs.images
     image_paths = [save_image(img) for img in images]
     return image_paths, seed