Update app.py

app.py

@@ -1,32 +1,79 @@
-import gradio as gr
-import torch
-import time
+import sys
+from unittest.mock import MagicMock
 
+import torch
 if not hasattr(torch, 'float8_e8m0fnu'):
-    torch.float8_e8m0fnu = torch.float16 
+    torch.float8_e8m0fnu = torch.float16
+
+import io
+from PIL import Image
+from huggingface_hub import InferenceClient
+
+class APIPipelineMock:
+    def __init__(self):
+        self.client = InferenceClient()
+        self.model_id = "simianluo/lcm_dreamshaper_v7"
+        self.scheduler = MagicMock()
+        
+    def to(self, device): return self
+    def enable_attention_slicing(self): pass
+    def enable_vae_slicing(self): pass
+    
+    def __call__(self, prompt, image=None, strength=0.5, **kwargs):
+        if image is None:
+            img = self.client.text_to_image(prompt, model=self.model_id)
+        else:
+            img_byte_arr = io.BytesIO()
+            image.save(img_byte_arr, format='PNG')
+            img_bytes = img_byte_arr.getvalue()
+            img = self.client.image_to_image(img_bytes, prompt=prompt, model=self.model_id, strength=strength)
+        
+        class DiffusersOutput:
+            def __init__(self, images): self.images = images
+        return DiffusersOutput(images=img)
+
+mock_pipe = APIPipelineMock()
+
+class MockDiffusionPipeline:
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs): return mock_pipe
 
+class MockAutoPipeline:
+    @classmethod
+    def from_pipe(cls, *args, **kwargs): return mock_pipe
+
+
+import diffusers
+diffusers.DiffusionPipeline = MockDiffusionPipeline
+diffusers.AutoPipelineForImage2Image = MockAutoPipeline
+
+
+import gradio as gr
+import time
 from LLM_pipeline import smart_generate
 from model_loading import GenerationSession
 
 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}")
+session.txt2img_pipeline = mock_pipe
+session.img2img_pipeline = mock_pipe
 
 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")
+    
+    if session.current_image is None:
+        raw_output = mock_pipe(prompt=user_prompt)
+        final_image = raw_output.images
+    else:
+        raw_output = mock_pipe(prompt=user_prompt, image=session.current_image, strength=0.45)
+        final_image = raw_output.images
         
+    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():
@@ -42,7 +89,7 @@ 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)
         time_output = gr.Textbox(label="Generation Time", interactive=False)