bugfix cursor gpu

app.py

@@ -11,24 +11,23 @@ from huggingface_hub import login
 hf_token = os.getenv("hf_token")
 login(token=hf_token)
 
-
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
 _text_gen_pipeline = None
+@spaces.GPU()
 def get_text_gen_pipeline():
     global _text_gen_pipeline
     if _text_gen_pipeline is None:
          try:
+              device = "cuda" if torch.cuda.is_available() else "cpu"
               _text_gen_pipeline = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.3", max_new_tokens=2048, device=device)
          except Exception as e:
               print(f"Error loading text generation model: {e}")
               return None
     return _text_gen_pipeline
 
+@spaces.GPU()
 def refine_prompt(prompt):
     text_gen = get_text_gen_pipeline()
     if text_gen is None:
@@ -53,6 +52,10 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
     try:
         progress(0, desc="Starting generation...")
         
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        dtype = torch.bfloat16
+        pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+        
         # Validate that prompt is not empty
         if not prompt or prompt.strip() == "":
             return None, "Please provide a valid prompt."