Upload 4 files

Add application files

LLM_pipeline.py

@@ -0,0 +1,31 @@
+import os
+import requests
+from model_loading import GenerationSession
+
+def prompt_enhancer(user_prompt: str) -> str:
+    try:
+        response = requests.post(
+            "http://localhost:11434/api/generate",
+            json={
+                "model": "mistral",
+                "format": "json",  
+                "prompt": (
+                    f"[INST] You are an image generation prompt engineer. "
+                    f"Rewrite this prompt to be vivid and detailed, under 60 words. "
+                    f"Return ONLY the rewritten prompt, nothing else.\n\n"
+                    f"Prompt: {user_prompt} [/INST]"
+                ),
+                "stream": False
+            },
+            timeout=60
+        )
+    except requests.exceptions.ConnectionError:
+        print("Warning: Could not connect to local Ollama.")
+        return user_prompt
+    return response.json()["response"].strip()
+
+def smart_generate(user_prompt: str, session: GenerationSession, strength: float = 0.45):
+    enhanced = prompt_enhancer(user_prompt)
+
+    image = session.Generate(enhanced, strength=strength)
+    return image[0], enhanced 

app.py

@@ -0,0 +1,41 @@
+import gradio as gr
+import torch
+if not hasattr(torch, 'float8_e8m0fnu'):
+    torch.float8_e8m0fnu = torch.float16 
+from LLM_pipeline import smart_generate
+from model_loading import GenerationSession
+import time
+
+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")
+    
+    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. Next generation will be a brand new Base Image."
+
+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.")
+    
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Describe the image you want to create or modify...")
+        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")
+        enhanced_prompt = gr.Textbox(label="Enhanced Prompt", interactive=False)
+
+    generate_button.click(fn=ui_handler, inputs=prompt_input, outputs=[output_image, enhanced_prompt, gr.Textbox(label="Generation Time", interactive=False)])
+    reset_button.click(fn=ui_reset, inputs=None, outputs=[output_image, enhanced_prompt, gr.Textbox(label="Generation Time", interactive=False)])
+
+if __name__ == "__main__":
+    demo.launch()

model_loading.py

@@ -0,0 +1,83 @@
+from PIL import Image
+import os
+import ollama
+import torch
+from diffusers import DiffusionPipeline, AutoPipelineForImage2Image, LCMScheduler, AutoPipelineForText2Image
+import time
+
+model_id = "simianluo/lcm_dreamshaper_v7"
+
+class GenerationSession:
+    def __init__(self, model_id):
+        self.model_id = model_id   
+        self.txt2img_pipeline = None  
+        self.img2img_pipeline = None
+        self.current_image = None
+        self.current_prompt = None
+        self._initialize_pipelines()
+
+    def _initialize_pipelines(self):
+            print("initializing pipelines...")
+
+            self.txt2img_pipeline = DiffusionPipeline.from_pretrained(
+                model_id,
+                torch_dtype = torch.float16,
+                safety_checker = None
+            )
+            
+            self.txt2img_pipeline.scheduler = LCMScheduler.from_config(self.txt2img_pipeline.scheduler.config)
+            self.txt2img_pipeline.to("cuda")
+            self.txt2img_pipeline.enable_attention_slicing()
+            self.txt2img_pipeline.enable_vae_slicing()
+
+           # self.txt2img_pipeline.unet = torch.compile(
+           #   self.txt2img_pipeline.unet,
+           #    mode = "reduce-overhead",
+           #     fullgraph = True
+           #)
+            print("Text 2 image pipeline loaded and compiled.")
+
+
+            self.img2img_pipeline = AutoPipelineForImage2Image.from_pipe(self.txt2img_pipeline)  
+            print("Image 2 image pipeline loaded (shared weights).")
+
+    def GeneratingBaseImage(self, prompt: str, negative_prompt: str = "Blurry, low quality, static and distorted image") -> str:
+        start = time.time()
+        image = self.txt2img_pipeline(
+            prompt = prompt,
+            negative_prompt= negative_prompt,
+            num_inference_steps = 4,
+            guidance_scale = 1.0,
+            height = 512,
+            width = 512
+        ).images
+        print(f"Text to image generated in [{time.time() - start:.2f}s]")
+        return image
+    
+    def GeneratingVariationImage(self, prompt: str, reference_image: Image.Image, strength: float = 0.5, negative_prompt: str = "Blurry, low quality, static and distorted image") -> str:
+        start = time.time()
+        image = self.img2img_pipeline(
+            prompt = prompt,
+            image = reference_image,
+            strength = strength,
+            num_inference_steps = 4,
+            guidance_scale = 1.0,
+            negative_prompt = negative_prompt
+        ).images
+        print(f"Image to image generated in [{time.time() - start:.2f}s]")
+        return image
+    
+    def Generate(self, new_prompt: str, strength: float = 0.5):
+        if self.current_image is None:
+            self.current_image = self.GeneratingBaseImage(new_prompt)
+        else:
+            self.current_image = self.GeneratingVariationImage(new_prompt, self.current_image, strength)
+        
+        self.current_prompt = new_prompt
+        return self.current_image
+    
+    def reset(self):
+        self.current_image = None
+        self.current_prompt = None
+        print("Session reset. Ready for new generation.")
+

requirements.txt

@@ -0,0 +1,14 @@
+# Core UI Framework
+gradio>=4.0.0
+
+# Deep Learning & Model Inference Frameworks
+torch
+torchvision
+transformers
+diffusers>=0.25.0
+accelerate>=0.26.0
+
+# Image Processing and Utilities
+Pillow
+numpy
+huggingface_hub