Update app.py

app.py

@@ -25,30 +25,26 @@ def dynamic_cloud_call(prompt, image=None, strength=0.45, **kwargs):
         image.save(img_byte_arr, format='PNG')
         img_bytes = img_byte_arr.getvalue()
         output_image = api_client.image_to_image(img_bytes, prompt=prompt, model=target_model, strength=strength)
-    
-    class DiffusersOutput:
-        def __init__(self, images):
-            self.images = [images] if not isinstance(images, list) else images
-    return DiffusersOutput(images=output_image)
-
-if session.txt2img_pipeline is not None:
-    session.txt2img_pipeline.__class__.__call__ = dynamic_cloud_call
+    return output_image
 
 def ui_handler(user_prompt):
     start_time = time.time()
-    image_list, enhanced_text = smart_generate(user_prompt, session, strength=0.45)
     
-    if isinstance(image_list, list) and len(image_list) > 0:
-        final_image = image_list
+    if session.current_image is None:
+        final_image = dynamic_cloud_call(prompt=user_prompt)
     else:
-        final_image = image_list
+        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
         
     end_time = time.time()
     return final_image, enhanced_text, f"Total generation time: {end_time - start_time:.2f}s"
 
 def ui_reset():
     session.reset()
-    return None, "Session cleared. Next generation will be a brand new Base Image.", "Session reset."
+    return None, "Session cleared.", "Session reset."
 
 with gr.Blocks(title="Active Image Generator", theme=gr.Theme.from_hub("Respair/Shiki")) as demo:
     gr.Markdown("## Active Image Generator\n\nEnter a prompt to generate or modify an image.")
@@ -71,9 +67,4 @@ with gr.Blocks(title="Active Image Generator", theme=gr.Theme.from_hub("Respair/
     
     reset_button.click(
         fn=ui_reset, 
-        inputs=None, 
-        outputs=[output_image, enhanced_prompt, time_output]
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+        inputs=None,