Update app.py

app.py

@@ -17,7 +17,7 @@ session = GenerationSession(model_id)
 api_client = InferenceClient()
 target_model = "simianluo/lcm_dreamshaper_v7"
 
-def dynamic_cloud_call(prompt, image=None, strength=0.45, **kwargs):
+def dynamic_cloud_call(prompt, image=None, strength=0.45):
     if image is None:
         output_image = api_client.text_to_image(prompt, model=target_model)
     else:
@@ -30,15 +30,27 @@ def dynamic_cloud_call(prompt, image=None, strength=0.45, **kwargs):
 def ui_handler(user_prompt):
     start_time = time.time()
     
-    if session.current_image is None:
-        final_image = dynamic_cloud_call(prompt=user_prompt)
-    else:
-        final_image = dynamic_cloud_call(prompt=user_prompt, image=session.current_image, strength=0.45)
-        
+    try:
+        image_list, enhanced_text = smart_generate(user_prompt, session, strength=0.45)
+        if isinstance(image_list, list) and len(image_list) > 0:
+            base_img = image_list
+        else:
+            base_img = image_list
+            
+        if hasattr(base_img, "save"):
+            final_image = base_img
+        else:
+            raise ValueError()
+    except Exception:
+        enhanced_text = user_prompt
+        if session.current_image is None:
+            final_image = dynamic_cloud_call(prompt=user_prompt)
+        else:
+            final_image = dynamic_cloud_call(prompt=user_prompt, image=session.current_image, strength=0.45)
+            
     session.current_image = final_image
-    session.current_prompt = user_prompt
-    enhanced_text = user_prompt
-        
+    session.current_prompt = enhanced_text
+    
     end_time = time.time()
     return final_image, enhanced_text, f"Total generation time: {end_time - start_time:.2f}s"
 
@@ -67,4 +79,9 @@ with gr.Blocks(title="Active Image Generator", theme=gr.Theme.from_hub("Respair/
     
     reset_button.click(
         fn=ui_reset, 
-        inputs=None,
+        inputs=None, 
+        outputs=[output_image, enhanced_prompt, time_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch()