From 68712fb418fbbe5d24915fd515d0f8131d257710 Mon Sep 17 00:00:00 2001
From: James Ketrenos <james_git@ketrenos.com>
Date: Thu, 22 May 2025 16:08:03 -0700
Subject: [PATCH] Image gen is working

---
 frontend/src/NewApp/Components/Pulse.tsx      |  6 ++---
 .../NewApp/Components/streamQueryResponse.tsx |  3 ++-
 .../src/NewApp/Pages/GenerateCandidate.tsx    | 22 +++++++++++--------
 src/utils/agents/image_generator.py           | 18 ++++++++++-----
 src/utils/profile_image.py                    | 15 ++++++++-----
 5 files changed, 39 insertions(+), 25 deletions(-)

diff --git a/frontend/src/NewApp/Components/Pulse.tsx b/frontend/src/NewApp/Components/Pulse.tsx
index db01dcb..54264f4 100644
--- a/frontend/src/NewApp/Components/Pulse.tsx
+++ b/frontend/src/NewApp/Components/Pulse.tsx
@@ -37,8 +37,8 @@ const Pulse: React.FC<PulseProps> = ({ timestamp, sx }) => {
   };
 
   const baseCoreStyle: React.CSSProperties = {
-    width: 24,
-    height: 24,
+      width: 0,
+      height: 0,
     borderRadius: '50%',
     backgroundColor: '#2196f3',
     position: 'relative',
@@ -92,8 +92,6 @@ const Pulse: React.FC<PulseProps> = ({ timestamp, sx }) => {
     animation: 'ripple-expand 1s ease-out 0.3s forwards',
   };
 
-
-
   return (
     <>
       <style>
diff --git a/frontend/src/NewApp/Components/streamQueryResponse.tsx b/frontend/src/NewApp/Components/streamQueryResponse.tsx
index d8711a8..61545aa 100644
--- a/frontend/src/NewApp/Components/streamQueryResponse.tsx
+++ b/frontend/src/NewApp/Components/streamQueryResponse.tsx
@@ -1,5 +1,6 @@
 import {  BackstoryMessage } from '../../Components/Message';
 import { Query } from '../../Components/ChatQuery';
+import { jsonrepair } from 'jsonrepair';
 
 type StreamQueryOptions = {
   query: Query;
@@ -66,7 +67,7 @@ const streamQueryResponse = (options: StreamQueryOptions) => {
       let streaming_response = '';
 
       const processLine = async (line: string) => {
-        const update = JSON.parse(line);
+        const update = JSON.parse(jsonrepair(line));
 
         switch (update.status) {
             case "streaming":
diff --git a/frontend/src/NewApp/Pages/GenerateCandidate.tsx b/frontend/src/NewApp/Pages/GenerateCandidate.tsx
index 7f8f138..5286922 100644
--- a/frontend/src/NewApp/Pages/GenerateCandidate.tsx
+++ b/frontend/src/NewApp/Pages/GenerateCandidate.tsx
@@ -64,7 +64,7 @@ const GenerateCandidate = (props: BackstoryElementProps) => {
 
         controllerRef.current = streamQueryResponse({
             query: {
-                prompt: `A photorealistic profile picture of a ${user.age} year old ${user.gender} ${user.ethnicity} person. ${prompt}`,
+                prompt: `A photorealistic profile picture of a ${userRef.current.age} year old ${userRef.current.gender} ${userRef.current.ethnicity} person. ${prompt}`,
                 agent_options: {
                     username: userRef.current.username,
                     filename: "profile.png"
@@ -74,7 +74,7 @@ const GenerateCandidate = (props: BackstoryElementProps) => {
             sessionId,
             connectionBase,
             onComplete: (msg) => {
-                console.log({ msg, state: stateRef.current, prompt: promptRef.current || '' });
+                console.log(msg);
                 switch (msg.status) {
                     case "partial":
                     case "done":
@@ -98,6 +98,9 @@ const GenerateCandidate = (props: BackstoryElementProps) => {
                         break;
                     default:
                         const data = JSON.parse(msg.response || '');
+                        if (msg.status !== "heartbeat") {
+                            console.log(data);
+                        }
                         if (data.timestamp) {
                             setTimestamp(data.timestamp);
                         } else {
@@ -156,6 +159,7 @@ const GenerateCandidate = (props: BackstoryElementProps) => {
                         if (msg.status === "done") {
                             setProcessing(false);
                             setCanGenImage(true);
+                            setStatus('');
                             controllerRef.current = null;
                             stateRef.current = 0;
                             setTimeout(() => {
@@ -263,17 +267,17 @@ const GenerateCandidate = (props: BackstoryElementProps) => {
             </Box>
         }
         <Box sx={{display: "flex", flexDirection: "column"}}>
-            <Box sx={{ display: "flex", flexDirection: "row"}}>
+                <Box sx={{ display: "flex", flexDirection: "row", position: "relative" }}>
                 <Avatar
-                    src={hasProfile ? `/api/u/${user.username}/profile/${sessionId}` : ''}
-                    alt={`${user.full_name}'s profile`}
+                        src={hasProfile ? `/api/u/${userRef.current.username}/profile/${sessionId}` : ''}
+                        alt={`${userRef.current.full_name}'s profile`}
                     sx={{
-                    width: 80,
-                    height: 80,
-                    border: '2px solid #e0e0e0',
+                        width: 80,
+                        height: 80,
+                        border: '2px solid #e0e0e0',
                     }}
                 />
-                { processing && <Pulse sx={{position: "absolute", left: "-80px" }} timestamp={timestamp}/> }
+                    {processing && <Pulse sx={{ position: "relative", left: "-80px" }} timestamp={timestamp} />}
                 <Tooltip title={"Generate Profile Picture"}>
                     <span style={{ display: "flex", flexGrow: 1 }}>
                     <Button
diff --git a/src/utils/agents/image_generator.py b/src/utils/agents/image_generator.py
index fe8d493..a3f2000 100644
--- a/src/utils/agents/image_generator.py
+++ b/src/utils/agents/image_generator.py
@@ -96,7 +96,9 @@ class ImageGenerator(Agent):
             logger.info("Beginning image generation...", self.user)
             logger.info("TODO: Add safety checks for filename... actually figure out an entirely different way to figure out where to store them.")
             self.filename = "profile.png"
-            request = ImageRequest(filepath=os.path.join(defines.user_dir, self.user["username"], self.filename), prompt=prompt)
+            file_path = os.path.join(defines.user_dir, self.user["username"], self.filename)
+            logger.info(f"Image generation: {file_path} <- {prompt}")
+            request = ImageRequest(filepath=file_path, prompt=prompt)
             async for message in generate_image(
                 message=message,
                 request=request 
@@ -104,11 +106,15 @@ class ImageGenerator(Agent):
                 if message.status != "done":
                     yield message
             logger.info("Image generation done...")
-            images = self.user.get("images", [])
-            if self.filename not in images:
-                images.append(self.filename)
-            if self.filename == "profile.png":
-                self.user["has_profile"] = True
+            if not os.path.exists(file_path):
+                logger.info(f"Generated image does not exist: {file_path}")
+                logger.error(f"{message.status} {message.response}")
+            else:
+                images = self.user.get("images", [])
+                if self.filename not in images:
+                    images.append(self.filename)
+                if self.filename == "profile.png":
+                    self.user["has_profile"] = True
 
             #
             # Write out the completed user information
diff --git a/src/utils/profile_image.py b/src/utils/profile_image.py
index a22747e..aecf6bf 100644
--- a/src/utils/profile_image.py
+++ b/src/utils/profile_image.py
@@ -72,14 +72,14 @@ def flux_worker(pipe: Any, params: ImageRequest, status_queue: queue.Queue, task
         # Simulate your pipe call with progress updates
         def status_callback(pipeline, step, timestep, callback_kwargs):
             # Send progress updates
-            elapsed = time.time() - start_gen_time
-            progress = int((step + 1) / params.iterations * 100)
+            progress = int((step+1) / params.iterations * 100)
             
             status_queue.put({
                 "status": "running",
-                "message": f"Processing step {step}/{params.iterations}",
+                "message": f"Processing step {step+1}/{params.iterations} ({progress}%)",
                 "progress": progress
             })
+            return callback_kwargs
         
         # Replace this block with your actual Flux pipe call:
         image = pipe(
@@ -91,10 +91,10 @@ def flux_worker(pipe: Any, params: ImageRequest, status_queue: queue.Queue, task
             callback_on_step_end=status_callback,
         ).images[0]
         
-        # Simulate image generation completion
         gen_time = time.time() - start_gen_time
         per_step_time = gen_time / params.iterations if params.iterations > 0 else gen_time
         
+        logger.info(f"Saving to {params.filepath}")
         image.save(params.filepath)
 
         # Final completion status
@@ -104,10 +104,12 @@ def flux_worker(pipe: Any, params: ImageRequest, status_queue: queue.Queue, task
             "progress": 100,
             "generation_time": gen_time,
             "per_step_time": per_step_time,
-            "image_path": f"generated_{task_id}.png"  # Replace with actual image path/URL
+            "image_path": params.filepath
         })
         
     except Exception as e:
+        logger.error(traceback.format_exc())
+        logger.error(e)
         status_queue.put({
             "status": "error",
             "message": f"Generation failed: {str(e)}",
@@ -198,6 +200,7 @@ async def async_generate_image(pipe: Any, params: ImageRequest) -> AsyncGenerato
             'message': f'Server error: {str(e)}',
             'error': str(e)
         }
+        logger.error(error_status)
         yield error_status
     
     finally:
@@ -253,6 +256,8 @@ async def generate_image(message: Message, request: ImageRequest) -> AsyncGenera
         async for update in async_generate_image(pipe, request):
             message.status = update.get("status", "thinking")
             message.response = json.dumps(update) # Merge properties from async_generate_image over the message...
+            if message.status != "heartbeat":
+                logger.info(message.response)
             yield message
         
         # Final result