Update app.py

app.py

@@ -1,28 +1,52 @@
 import gradio as gr
 import torch
 import time
+import io
+from PIL import Image
+from huggingface_hub import InferenceClient
 
-if not hasattr(torch, 'float8_e8m0fnu'):
-    torch.float8_e8m0fnu = torch.float16 
 
-from LLM_pipeline import smart_generate
-from model_loading import GenerationSession
+class CloudFastSession:
+    def __init__(self):
+        self.client = InferenceClient()
+        self.model_id = "simianluo/lcm_dreamshaper_v7"
+        self.current_image = None
+        self.current_prompt = None
 
-model_id = "runwayml/stable-diffusion-v1-5"
-session = GenerationSession(model_id)
-
-def ui_handler(user_prompt):
-    start_time = time.time()
-    image_list, enhanced_text = smart_generate(user_prompt, session, strength=0.45)
-    final_image = image_list if isinstance(image_list, list) else image_list
-    end_time = time.time()
-    print(f"Image generation time: {end_time:.2f}s")
+    def ui_handler(self, user_prompt):
+        start_time = time.time()
         
-    return final_image, enhanced_text, f"Total generation time: {end_time - start_time:.2f}s"
+        try:
+            if self.current_image is None:
+                image = self.client.text_to_image(user_prompt, model=self.model_id)
+            else:
+                img_byte_arr = io.BytesIO()
+                self.current_image.save(img_byte_arr, format='PNG')
+                img_bytes = img_byte_arr.getvalue()
+                
+                image = self.client.image_to_image(img_bytes, prompt=user_prompt, model=self.model_id, strength=0.45)
+            
+            self.current_image = image
+            self.current_prompt = user_prompt
+            enhanced_text = user_prompt
+            
+        except Exception as e:
+            print(f"API Error: {e}. Falling back to a clean text-to-image call.")
+            image = self.client.text_to_image(user_prompt, model=self.model_id)
+            self.current_image = image
+            self.current_prompt = user_prompt
+            enhanced_text = user_prompt
+
+        end_time = time.time()
+        return self.current_image, enhanced_text, f"Total generation time: {end_time - start_time:.2f}s"
+
+    def ui_reset(self):
+        self.current_image = None
+        self.current_prompt = None
+        return None, "Session cleared.", "Session reset. Next generation will be a brand new Base Image."
+
 
-def ui_reset():
-    session.reset()
-    return None, "Session cleared. Next generation will be a brand new Base Image.", "Session reset. Next generation will be a brand new Base Image."
+session_manager = CloudFastSession()
 
 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. Each new prompt will build upon the previous image, creating a dynamic and evolving visual experience. Use the reset button to start fresh with a new base image.")
@@ -33,20 +57,18 @@ with gr.Blocks(title="Active Image Generator", theme=gr.Theme.from_hub("Respair/
         reset_button = gr.Button("Reset Session", variant="secondary")
         
     with gr.Column():
-        output_image = gr.Image(label="Generated Image")
+        output_image = gr.Image(label="Generated Image", type="pil")
         enhanced_prompt = gr.Textbox(label="Enhanced Prompt", interactive=False)
-        # Fix: Defined the component here and assigned it to a variable
         time_output = gr.Textbox(label="Generation Time", interactive=False)
 
-    # Fix: Used the variable 'time_output' in the outputs list
     generate_button.click(
-        fn=ui_handler, 
+        fn=session_manager.ui_handler, 
         inputs=prompt_input, 
         outputs=[output_image, enhanced_prompt, time_output]
     )
     
     reset_button.click(
-        fn=ui_reset, 
+        fn=session_manager.ui_reset, 
         inputs=None, 
         outputs=[output_image, enhanced_prompt, time_output]
     )