simplify code

app.py

@@ -3,14 +3,14 @@ import numpy as np
 import random
 import spaces
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import FluxPipeline
 
+# Check for CUDA and set device
 device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
 
-# Load the model in FP16
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.float16)
-
-# Move the pipeline to GPU if available
+# Load the model
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.float16)
 pipe = pipe.to(device)
 
 # Convert text encoders to full precision
@@ -18,48 +18,36 @@ pipe.text_encoder = pipe.text_encoder.to(torch.float32)
 if hasattr(pipe, 'text_encoder_2'):
     pipe.text_encoder_2 = pipe.text_encoder_2.to(torch.float32)
 
-# Enable memory efficient attention if available and on CUDA
-if device == "cuda" and hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
-    try:
-        pipe.enable_xformers_memory_efficient_attention()
-        print("xformers memory efficient attention enabled")
-    except Exception as e:
-        print(f"Could not enable memory efficient attention: {e}")
-
-# Compile the UNet for potential speedups if on CUDA
-if device == "cuda":
-    try:
-        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        print("UNet compiled for potential speedups")
-    except Exception as e:
-        print(f"Could not compile UNet: {e}")
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
 @spaces.GPU()
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    # Use full precision for text encoding
-    with torch.no_grad():
-        text_inputs = pipe.tokenizer(prompt, return_tensors="pt").to(device)
-        text_embeddings = pipe.text_encoder(text_inputs.input_ids)[0]
-    
-    # Use mixed precision for the rest of the pipeline
-    with torch.inference_mode(), torch.autocast(device_type='cuda', dtype=torch.float16):
-        image = pipe(
-            prompt_embeds=text_embeddings,
-            width=width,
-            height=height,
-            num_inference_steps=num_inference_steps,
-            generator=generator,
-            guidance_scale=0.0
-        ).images[0]
-    
-    return image, seed
+    try:
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+        generator = torch.Generator(device=device).manual_seed(seed)
+        
+        # Use full precision for text encoding
+        with torch.no_grad():
+            text_inputs = pipe.tokenizer(prompt, return_tensors="pt").to(device)
+            text_embeddings = pipe.text_encoder(text_inputs.input_ids)[0]
+        
+        # Use mixed precision for the rest of the pipeline
+        with torch.autocast(device_type=device, dtype=torch.float16):
+            image = pipe(
+                prompt_embeds=text_embeddings,
+                width=width,
+                height=height,
+                num_inference_steps=num_inference_steps,
+                generator=generator,
+                guidance_scale=0.0
+            ).images[0]
+        
+        return image, seed
+    except Exception as e:
+        print(f"Error during inference: {e}")
+        return None, seed
  
 examples = [
     "a tiny astronaut hatching from an egg on the moon",