Update app.py

app.py

@@ -64,9 +64,10 @@ def generate(input_ids, width, height, cfg_weight=5, temperature=1.0, parallel_s
             generated_tokens = torch.zeros((parallel_size, 576), dtype=torch.int, device=device)
 
             pkv = None
-            for i in range(576):
-                if progress:
-                    progress((i + 1) / 576, desc="Generating image tokens")
+            total_steps = 576
+            for i in range(total_steps):
+                if progress is not None:
+                    progress((i + 1) / total_steps, desc="Generating image tokens")
                 
                 outputs = vl_gpt.language_model.model(
                     inputs_embeds=inputs_embeds,
@@ -114,7 +115,8 @@ def generate_image(prompt, seed=None, guidance=5, t2i_temperature=1.0, progress=
         if not prompt.strip():
             raise gr.Error("Please enter a valid prompt.")
 
-        progress(0, desc="Initializing...")
+        if progress is not None:
+            progress(0, desc="Initializing...")
         torch.cuda.empty_cache()
 
         # Seed management
@@ -135,7 +137,9 @@ def generate_image(prompt, seed=None, guidance=5, t2i_temperature=1.0, progress=
         ) + vl_chat_processor.image_start_tag
 
         input_ids = torch.tensor(tokenizer.encode(text), dtype=torch.long, device=device)
-        progress(0.1, desc="Generating image tokens...")
+        
+        if progress is not None:
+            progress(0.1, desc="Generating image tokens...")
         
         generated_tokens = generate(
             input_ids,
@@ -147,7 +151,8 @@ def generate_image(prompt, seed=None, guidance=5, t2i_temperature=1.0, progress=
             progress=progress
         )
         
-        progress(0.9, desc="Processing images...")
+        if progress is not None:
+            progress(0.9, desc="Processing images...")
         patches = vl_gpt.gen_vision_model.decode_code(
             generated_tokens.to(dtype=torch.int),
             shape=[PARALLEL_SIZE, 8, DEFAULT_WIDTH // PATCH_SIZE, DEFAULT_HEIGHT // PATCH_SIZE]
@@ -157,8 +162,11 @@ def generate_image(prompt, seed=None, guidance=5, t2i_temperature=1.0, progress=
         return images
 
     except Exception as e:
-        logger.error(f"Generation failed: {str(e)}")
-        raise gr.Error(f"Image generation failed: {str(e)}")
+        logger.error(f"Generation failed: {str(e)}", exc_info=True)
+        if "index out of range" in str(e).lower():
+            raise gr.Error("Image generation failed due to internal error. Please try again with different parameters.")
+        else:
+            raise gr.Error(f"Image generation failed: {str(e)}")
 
 def create_interface():
     with gr.Blocks(title="Janus-Pro-7B Image Generator", theme=gr.themes.Soft()) as demo:
@@ -169,7 +177,11 @@ def create_interface():
 
         with gr.Row():
             with gr.Column(scale=3):
-                prompt_input = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate...", lines=3)
+                prompt_input = gr.Textbox(
+                    label="Prompt", 
+                    placeholder="Describe the image you want to generate...", 
+                    lines=3
+                )
                 generate_btn = gr.Button("Generate Images", variant="primary")
                 
                 with gr.Accordion("Advanced Settings", open=False):
@@ -178,22 +190,36 @@ def create_interface():
                             label="Seed", 
                             value=None, 
                             precision=0, 
-                            info="Leave empty for random seed"  # Fixed parameter
+                            info="Leave empty for random seed"
                         )
                         guidance_slider = gr.Slider(
-                            3, 10, value=5, step=0.5,
                             label="CFG Guidance Weight",
+                            minimum=3,
+                            maximum=10,
+                            value=5,
+                            step=0.5,
                             info="Higher values = more prompt adherence, lower values = more creativity"
                         )
                         temp_slider = gr.Slider(
-                            0.1, 1.0, value=1.0, step=0.1,
                             label="Temperature",
+                            minimum=0.1,
+                            maximum=1.0,
+                            value=1.0,
+                            step=0.1,
                             info="Higher values = more randomness, lower values = more deterministic"
                         )
 
             with gr.Column(scale=2):
-                output_gallery = gr.Gallery(label="Generated Images", columns=2, height=600, preview=True)
-                status = gr.Textbox(label="Status", interactive=False)
+                output_gallery = gr.Gallery(
+                    label="Generated Images", 
+                    columns=2, 
+                    height=600, 
+                    preview=True
+                )
+                status = gr.Textbox(
+                    label="Status", 
+                    interactive=False
+                )
 
         gr.Examples(
             examples=[