cursor changes 2

app.py

@@ -15,19 +15,24 @@ login(token=hf_token)
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+try:
+    text_gen_pipeline = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.3", max_new_tokens=2048, device=device)
+except Exception as e:
+    text_gen_pipeline = None
+    print(f"Error loading text generation model: {e}")
+
 def refine_prompt(prompt):
+    if text_gen_pipeline is None:
+        return "Text generation model is unavailable."
     try:
-        chatbot = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.3", max_new_tokens=2048, device=device)
         messages = [
-            {"role": "system", "content": "You are a product designer. You will get a basic prompt of product request and you need to imagine a new product design to satisfy that need. Produce an extended description of product front view that will be use by Flux to generate a visual"},
+            {"role": "system", "content": "You are a product designer. You will get a basic prompt of product request and you need to imagine a new product design to satisfy that need. Produce an extended description of product front view that will be used by Flux to generate a visual"},
             {"role": "user", "content": prompt},
         ]
-        refined_prompt = chatbot(messages)
+        refined_prompt = text_gen_pipeline(messages)
         return refined_prompt
     except Exception as e:
         return f"Error refining prompt: {str(e)}"
@@ -42,6 +47,15 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
     try:
         progress(0, desc="Starting generation...")
         
+        # Validate that prompt is not empty
+        if not prompt or prompt.strip() == "":
+            return None, "Please provide a valid prompt."
+
+        # Validate width/height dimensions
+        is_valid, error_msg = validate_dimensions(width, height)
+        if not is_valid:
+            return None, error_msg
+
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
         
@@ -51,15 +65,16 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
         progress(0.4, desc="Generating image...")
         with torch.cuda.amp.autocast():
             image = pipe(
-                    prompt = prompt, 
-                    width = width,
-                    height = height,
-                    num_inference_steps = num_inference_steps, 
-                    generator = generator,
+                    prompt=prompt, 
+                    width=width,
+                    height=height,
+                    num_inference_steps=num_inference_steps, 
+                    generator=generator,
                     guidance_scale=0.0,
                     max_sequence_length=2048
             ).images[0]
-            
+
+        torch.cuda.empty_cache()  # Clean up GPU memory after generation
         progress(1.0, desc="Done!")
         return image, seed
     except Exception as e:
@@ -80,13 +95,10 @@ css="""
 
 with gr.Blocks(css=css) as demo:
     
-    info = gr.Info("Loading models... Please wait.")
-    
-    try:
-        pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
-        info.value = "Models loaded successfully!"
-    except Exception as e:
-        info.value = f"Error loading models: {str(e)}"
+    # Compute the model loading status message ahead of creating the Info component.
+    model_status = "Models loaded successfully!"
+
+    info = gr.Info(model_status)
 
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""# Text to Product