Update utils.py

utils.py

@@ -1,24 +1,20 @@
-
 from diffusers import StableDiffusionPipeline
 import torch
+import os
+import requests
 
-# Load the model once globally
+# Initialize model only once at the top
 model_id = "stabilityai/stable-diffusion-2"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipe = pipe.to("cuda")
+pipe = StableDiffusionPipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+)
+pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-#def generate_image_from_prompt(prompt):
-#    image = pipe(prompt).images[0]
-#    return image
 def generate_image_from_prompt(prompt):
-    styled_prompt = f"{prompt}, in 2D cartoon style"
-    image = pipe(styled_prompt).images[0]
+    image = pipe(prompt).images[0]
     return image
 
-# Optional: GROQ function (if needed)
-import os
-import requests
-
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 
 def split_into_scenes(story_text, num_scenes=5):
@@ -35,6 +31,7 @@ Respond in a numbered list format.
         "Authorization": f"Bearer {GROQ_API_KEY}",
         "Content-Type": "application/json"
     }
+
     json_data = {
         "messages": [
             {"role": "system", "content": "You are a screenplay breakdown assistant."},
@@ -50,6 +47,7 @@ Respond in a numbered list format.
         return []
 
     output = response.json()["choices"][0]["message"]["content"]
+
     scenes = []
     for line in output.strip().split('\n'):
         line = line.strip()