Update app.py

app.py

@@ -1,15 +1,18 @@
 import spaces  # Import this first to avoid CUDA initialization issues
 import os
 import gradio as gr
+import numpy as np
 import json
 import torch
 import random
 import time
 from PIL import Image
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel
+from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
 
 # Define the device
-device = "cuda"
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Use the 'waffles' environment variable as the access token
 hf_token = os.getenv('waffles')
@@ -24,13 +27,10 @@ with open('loras.json', 'r') as f:
 
 # Initialize the base model with authentication and specify the device
 # Initialize the base model with authentication and specify the device
-pipe = DiffusionPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-schnell",
-    torch_dtype=torch.bfloat16,
-    token=hf_token
-).to(device)
+pipe = DiffusionPipeline.from_pretrained("sayakpaul/FLUX.1-merged", torch_dtype=dtype, token=hf_token).to(device)
 
-MAX_SEED = 2**32 - 1
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
 
 class calculateDuration:
     def __init__(self, activity_name=""):