try fix flux based on example

app.py

@@ -28,11 +28,12 @@ def get_image_gen_pipeline():
             dtype = torch.bfloat16
             _image_gen_pipeline = DiffusionPipeline.from_pretrained(
                 "black-forest-labs/FLUX.1-schnell",
-                # "black-forest-labs/FLUX.1-dev",
                 torch_dtype=dtype,
             ).to(device)
-            _image_gen_pipeline.enable_model_cpu_offload()
-            _image_gen_pipeline.enable_vae_slicing()
+            
+            # Comment these out for now to match the working example
+            # _image_gen_pipeline.enable_model_cpu_offload()
+            # _image_gen_pipeline.enable_vae_slicing()
         except Exception as e:
             print(f"Error loading image generation model: {e}")
             return None
@@ -94,9 +95,8 @@ def validate_dimensions(width, height):
     return True, None
 
 @spaces.GPU()
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4): 
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)): 
     try:
-        
         # Validate that prompt is not empty
         if not prompt or prompt.strip() == "":
             return None, "Please provide a valid prompt."
@@ -112,22 +112,20 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
         
-        generator = torch.Generator("cuda").manual_seed(seed)  # Explicitly use CUDA generator
+        # Use default torch generator instead of cuda-specific generator
+        generator = torch.Generator().manual_seed(seed)
         
-       # with torch.autocast('cuda'):
+        # Match the working example's parameters
         output = pipe(
-                prompt=prompt, 
-                width=width,
-                height=height,
-                num_inference_steps=num_inference_steps, 
-                generator=generator,
-                guidance_scale=7.5,
-                max_sequence_length=512
-            )
+            prompt=prompt, 
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps, 
+            generator=generator,
+            guidance_scale=0.0,  # Changed from 7.5 to 0.0
+        )
             
         image = output.images[0]
-
-        #torch.cuda.empty_cache()  
         return image, f"Image generated successfully with seed {seed}"
     except Exception as e:
         print(f"Error in infer: {str(e)}")  

requirements.txt

@@ -30,4 +30,8 @@ accelerate
 git+https://github.com/huggingface/diffusers.git
 invisible_watermark
 sentencepiece
-protobuf
+protobuf
+
+torch
+transformers==4.42.4
+xformers