Initial Commit

app.py

@@ -177,37 +177,59 @@ class FluxInterface:
 
     def generate(self, input_image, prompt="", guidance_scale=3.5, num_inference_steps=28, num_images=2, seed=None, aspect_ratio="1:1"):
         try:
+            print(f"Starting generation with prompt: {prompt}, guidance_scale: {guidance_scale}, steps: {num_inference_steps}")
+            
+            if input_image is None:
+                raise ValueError("No input image provided")
+                
             if seed is not None:
                 torch.manual_seed(seed)
+                print(f"Set random seed to: {seed}")
                 
             self.load_models()
+            print("Models loaded successfully")
             
             # Get dimensions from aspect ratio
             if aspect_ratio not in ASPECT_RATIOS:
                 raise ValueError(f"Invalid aspect ratio. Choose from {list(ASPECT_RATIOS.keys())}")
             width, height = ASPECT_RATIOS[aspect_ratio]
+            print(f"Using dimensions: {width}x{height}")
             
             # Process input image
-            input_image = self.resize_image(input_image)
-            qwen2_hidden_state, image_grid_thw = self.process_image(input_image)
-            pooled_prompt_embeds = self.compute_text_embeddings("")
+            try:
+                input_image = self.resize_image(input_image)
+                print(f"Input image resized to: {input_image.size}")
+                qwen2_hidden_state, image_grid_thw = self.process_image(input_image)
+                print("Input image processed successfully")
+            except Exception as e:
+                raise RuntimeError(f"Error processing input image: {str(e)}")
             
-            # Get T5 embeddings if prompt is provided
-            t5_prompt_embeds = self.compute_t5_text_embeddings(prompt)
+            try:
+                pooled_prompt_embeds = self.compute_text_embeddings("")
+                print("Base text embeddings computed")
+                
+                # Get T5 embeddings if prompt is provided
+                t5_prompt_embeds = self.compute_t5_text_embeddings(prompt)
+                print("T5 prompt embeddings computed")
+            except Exception as e:
+                raise RuntimeError(f"Error computing embeddings: {str(e)}")
             
             # Generate images
-            output_images = self.pipeline(
-                prompt_embeds=qwen2_hidden_state.repeat(num_images, 1, 1),
-                pooled_prompt_embeds=pooled_prompt_embeds,
-                t5_prompt_embeds=t5_prompt_embeds.repeat(num_images, 1, 1) if t5_prompt_embeds is not None else None,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=guidance_scale,
-                height=height,
-                width=width,
-            ).images
-            
-            return output_images
-            
+            try:
+                output_images = self.pipeline(
+                    prompt_embeds=qwen2_hidden_state.repeat(num_images, 1, 1),
+                    pooled_prompt_embeds=pooled_prompt_embeds,
+                    t5_prompt_embeds=t5_prompt_embeds.repeat(num_images, 1, 1) if t5_prompt_embeds is not None else None,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=guidance_scale,
+                    height=height,
+                    width=width,
+                ).images
+                
+                print("Images generated successfully")
+                return output_images
+            except Exception as e:
+                raise RuntimeError(f"Error generating images: {str(e)}")
         except Exception as e:
             print(f"Error during generation: {str(e)}")
             raise gr.Error(f"Generation failed: {str(e)}")
@@ -327,6 +349,7 @@ with gr.Blocks(
                     allow_preview=True,
                     preview=True
                 )
+                error_message = gr.Textbox(visible=False)
         
         with gr.Row(elem_classes="footer"):
             gr.Markdown("""
@@ -339,8 +362,20 @@ with gr.Blocks(
             """)
     
     # Set up the generation function
+    def generate_with_error_handling(*args):
+        try:
+            with gr.Status() as status:
+                status.update(value="Loading models...", visible=True)
+                results = interface.generate(*args)
+                status.update(value="Generation complete!", visible=False)
+                return [results, None]
+        except Exception as e:
+            error_msg = str(e)
+            print(f"Error in generate_with_error_handling: {error_msg}")
+            return [None, gr.Error(error_msg)]
+
     submit_btn.click(
-        fn=interface.generate,
+        fn=generate_with_error_handling,
         inputs=[
             input_image,
             prompt,
@@ -350,7 +385,10 @@ with gr.Blocks(
             seed,
             aspect_ratio
         ],
-        outputs=output_gallery,
+        outputs=[
+            output_gallery,
+            error_message
+        ],
         show_progress="minimal"
     )