Update app.py

app.py

@@ -1,4 +1,3 @@
-import spaces  # Import this first to avoid CUDA initialization issues
 import os
 import gradio as gr
 import numpy as np
@@ -7,18 +6,13 @@ from accelerate import dispatch_model, infer_auto_device_map
 from accelerate.utils import get_balanced_memory
 from torch.cuda.amp import autocast
 import torch
+import spaces  # Import this first to avoid CUDA initialization issues
 import random
 import time
 from PIL import Image
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel
 from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
 
-
-
-# Define the device
-dtype = torch.bfloat16
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-
 # Use the 'waffles' environment variable as the access token
 hf_token = os.getenv('waffles')
 
@@ -26,6 +20,14 @@ hf_token = os.getenv('waffles')
 if not hf_token:
     raise ValueError("Hugging Face API token not found. Please set the 'waffles' environment variable.")
 
+# Define the device
+dtype = torch.bfloat16
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+count0 = torch.zeros(1).to(device)
+count1 = torch.zeros(1).to(device)
+count2 = torch.zeros(1).to(device)
+
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
     loras = json.load(f)