Update app.py

app.py

@@ -1,74 +1,59 @@
 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 
 
-class CloudFastSession:
-    def __init__(self):
-        self.client = InferenceClient()
-        self.model_id = "simianluo/lcm_dreamshaper_v7"
-        self.current_image = None
-        self.current_prompt = None
+from LLM_pipeline import smart_generate
+from model_loading import GenerationSession
 
-    def ui_handler(self, user_prompt):
-        start_time = time.time()
-        
-        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."
+model_id = "simianluo/lcm_dreamshaper_v7"
+session = GenerationSession(model_id)
 
+if session.device == "cpu":
+    print("Applying High-Quality CPU sequential offload patches...")
+    try:
+        
+        session.txt2img_pipeline.enable_sequential_cpu_offload()
+        session.img2img_pipeline.enable_sequential_cpu_offload()
+    except Exception as e:
+        print(f"Note: Multi-pipeline offload optimized via sub-modules: {e}")
+
+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")
+        
+    return final_image, enhanced_text, f"Total generation time: {end_time - start_time:.2f}s"
 
-session_manager = CloudFastSession()
+def ui_reset():
+    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. 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.")
+    gr.Markdown("## Active Image Generator\n\nEnter a prompt to generate or modify an image.")
         
     with gr.Row():
-        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Describe the image you want to create or modify...")
+        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Describe the image...")
         generate_button = gr.Button("Generate", variant="primary")
         reset_button = gr.Button("Reset Session", variant="secondary")
         
     with gr.Column():
-        output_image = gr.Image(label="Generated Image", type="pil")
+        output_image = gr.Image(label="Generated Image")
         enhanced_prompt = gr.Textbox(label="Enhanced Prompt", interactive=False)
         time_output = gr.Textbox(label="Generation Time", interactive=False)
 
     generate_button.click(
-        fn=session_manager.ui_handler, 
+        fn=ui_handler, 
         inputs=prompt_input, 
         outputs=[output_image, enhanced_prompt, time_output]
     )
     
     reset_button.click(
-        fn=session_manager.ui_reset, 
+        fn=ui_reset, 
         inputs=None, 
         outputs=[output_image, enhanced_prompt, time_output]
     )