inital deployment1

app.py

@@ -1,27 +1,13 @@
-import torch
-import gradio as gr
-from diffusers import StableDiffusionPipeline, LCMScheduler
-
-
 import gradio as gr
 import cv2
+import torch
 import numpy as np
+from diffusers import StableDiffusionImg2ImgPipeline
+from PIL import Image
 
-# Function to process the video frame
-
-
-# Load the pre-trained Real-Time LCM model
-model_id = "SimianLuo/LCM_Dreamshaper_v7"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-pipe.to("cuda")
-
-# Function to generate images
-def generate_image(prompt: str):
-    image = pipe(prompt, num_inference_steps=4).images[0]
-    return image
-
-
+# Load the Stable Diffusion Model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
 
 def apply_color_filter(frame, filter_type="None"):
     """Apply a color filter to the frame."""
@@ -36,25 +22,39 @@ def apply_color_filter(frame, filter_type="None"):
         frame[:, :, 1] = 0  # Remove green channel
     return frame
 
-def process_frame(frame, filter_type="None"):
-    """Process a single frame by applying a color filter."""
+def process_frame(frame, filter_type="None", prompt="A futuristic landscape"):
+    """Process a single frame by applying a color filter and diffusion model."""
     frame = apply_color_filter(frame, filter_type)
-    return frame
-
-def video_stream(frame):
-    """Receives video from webcam and processes it"""
-    if frame is None:
-        return None
     
-    frame = cv2.imdecode(np.frombuffer(frame, np.uint8), cv2.IMREAD_COLOR)
-    processed_frame = process_frame(frame)
+    # Convert frame to PIL image
+    image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).resize((512, 512))
     
-    return processed_frame
-
-# Gradio Interface
+    # Apply Stable Diffusion model
+    result = pipe(prompt=prompt, image=image, strength=0.5, guidance_scale=7.5).images[0]
+    return np.array(result)
+
+def video_stream(filter_type, prompt):
+    """Captures video feed from webcam, applies color filter, and sends to the AI model."""
+    cap = cv2.VideoCapture(0)
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        frame = process_frame(frame, filter_type, prompt)
+        yield frame  # Return processed frame
+    cap.release()
+
+# Create Gradio App
 iface = gr.Interface(
     fn=video_stream,
-    inputs=gr.Video(source="webcam", streaming=True),  # Connect local webcam
-    outputs=gr.Image(label="Processed Webcam Feed"),   # Display processed frames
+    inputs=[
+        gr.Radio(["None", "Red", "Green", "Blue"], label="Color Filter"),
+        gr.Textbox(label="Stable Diffusion Prompt", value="A futuristic landscape")
+    ],
+    outputs=gr.Image(label="AI-Enhanced Webcam Feed"),
     live=True
 )
+
+if __name__ == "__main__":
+    iface.launch(share=True)