debug cursor image

app.py

@@ -15,13 +15,37 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
 _text_gen_pipeline = None
+_image_gen_pipeline = None
+
+@spaces.GPU()
+def get_image_gen_pipeline():
+    global _image_gen_pipeline
+    if _image_gen_pipeline is None:
+        try:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            dtype = torch.bfloat16
+            _image_gen_pipeline = DiffusionPipeline.from_pretrained(
+                "black-forest-labs/FLUX.1-schnell",
+                torch_dtype=dtype
+            ).to(device)
+        except Exception as e:
+            print(f"Error loading image generation model: {e}")
+            return None
+    return _image_gen_pipeline
+
 @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)
+              _text_gen_pipeline = pipeline(
+                   "text-generation",
+                   model="mistralai/Mistral-7B-Instruct-v0.3",
+                   max_new_tokens=2048,
+                   device=device,
+                   tokenizer_kwargs={"add_prefix_space": False}
+              )
          except Exception as e:
               print(f"Error loading text generation model: {e}")
               return None
@@ -52,9 +76,9 @@ 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)
+        pipe = get_image_gen_pipeline()
+        if pipe is None:
+            return None, "Image generation model is unavailable."
         
         # Validate that prompt is not empty
         if not prompt or prompt.strip() == "":
@@ -79,7 +103,7 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
                     height=height,
                     num_inference_steps=num_inference_steps, 
                     generator=generator,
-                    guidance_scale=0.0,
+                    guidance_scale=5.0,
                     max_sequence_length=2048
             ).images[0]
 
@@ -138,7 +162,7 @@ with gr.Blocks(css=css) as demo:
         
         run_button = gr.Button("Create visual", scale=0)
 
-        result = gr.Image(label="Result", show_label=False)
+        generated_image = gr.Image(label="Generated Image", show_label=False)
         
         with gr.Accordion("Advanced Settings Mistral", open=False):
             gr.Slider(
@@ -223,7 +247,7 @@ with gr.Blocks(css=css) as demo:
             examples = examples,
             fn = infer,
             inputs = [prompt],
-            outputs = [result, seed],
+            outputs = [generated_image, seed],
             cache_examples="lazy"
         )
 
@@ -239,7 +263,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click],
         fn = infer,
         inputs = [refined_prompt, seed, randomize_seed, width, height, num_inference_steps],
-        outputs = [result, seed]
+        outputs = [generated_image, seed]
     )
 
 demo.launch()