Update app.py

app.py

@@ -1,90 +1,60 @@
-import sys
-from unittest.mock import MagicMock
-
+import gradio as gr
 import torch
-if not hasattr(torch, 'float8_e8m0fnu'):
-    torch.float8_e8m0fnu = torch.float16
-
+import time
 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
-
+if not hasattr(torch, 'float8_e8m0fnu'):
+    torch.float8_e8m0fnu = torch.float16 
 
-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)
 
-session.txt2img_pipeline = mock_pipe
-session.img2img_pipeline = mock_pipe
+api_client = InferenceClient()
+target_model = "simianluo/lcm_dreamshaper_v7"
+
+def dynamic_cloud_call(prompt, image=None, strength=0.45, **kwargs):
+    if image is None:
+        output_image = api_client.text_to_image(prompt, model=target_model)
+    else:
+        img_byte_arr = io.BytesIO()
+        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
 
 def ui_handler(user_prompt):
     start_time = time.time()
+    image_list, enhanced_text = smart_generate(user_prompt, session, strength=0.45)
     
-    if session.current_image is None:
-        raw_output = mock_pipe(prompt=user_prompt)
-        final_image = raw_output.images
+    if isinstance(image_list, list) and len(image_list) > 0:
+        final_image = image_list
     else:
-        raw_output = mock_pipe(prompt=user_prompt, image=session.current_image, strength=0.45)
-        final_image = raw_output.images
+        final_image = image_list
         
-    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.", "Session reset."
+    return None, "Session cleared. Next generation will be a brand new Base Image.", "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.")
         
     with gr.Row():
-        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Describe the image...")
+        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Describe the image you want to create...")
         generate_button = gr.Button("Generate", variant="primary")
         reset_button = gr.Button("Reset Session", variant="secondary")