cursor changes

app.py

@@ -21,29 +21,49 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
 def refine_prompt(prompt):
-    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": "user", "content": prompt},
-    ]
-    refined_prompt = chatbot(messages)
-    return refined_prompt
+    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": "user", "content": prompt},
+        ]
+        refined_prompt = chatbot(messages)
+        return refined_prompt
+    except Exception as e:
+        return f"Error refining prompt: {str(e)}"
+
+def validate_dimensions(width, height):
+    if width * height > MAX_IMAGE_SIZE * MAX_IMAGE_SIZE:
+        return False, "Image dimensions too large"
+    return True, None
 
 @spaces.GPU()
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
-    image = pipe(
-            prompt = prompt, 
-            width = width,
-            height = height,
-            num_inference_steps = num_inference_steps, 
-            generator = generator,
-            guidance_scale=0.0,
-            max_sequence_length=2048
-    ).images[0] 
-    return image, seed
+    try:
+        progress(0, desc="Starting generation...")
+        
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+        
+        progress(0.2, desc="Setting up generator...")
+        generator = torch.Generator().manual_seed(seed)
+        
+        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,
+                    guidance_scale=0.0,
+                    max_sequence_length=2048
+            ).images[0]
+            
+        progress(1.0, desc="Done!")
+        return image, seed
+    except Exception as e:
+        return None, f"Error generating image: {str(e)}"
  
 examples = [
     "a tiny astronaut hatching from an egg on the moon",
@@ -60,7 +80,13 @@ css="""
 
 with gr.Blocks(css=css) as demo:
     
-    info = gr.Info("...")
+    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)}"
 
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""# Text to Product