Image gen is working

This commit is contained in:
James Ketr 2025-05-22 16:08:03 -07:00
parent 38b5185cfd
commit 68712fb418
5 changed files with 39 additions and 25 deletions

View File

@ -37,8 +37,8 @@ const Pulse: React.FC<PulseProps> = ({ timestamp, sx }) => {
}; };
const baseCoreStyle: React.CSSProperties = { const baseCoreStyle: React.CSSProperties = {
width: 24, width: 0,
height: 24, height: 0,
borderRadius: '50%', borderRadius: '50%',
backgroundColor: '#2196f3', backgroundColor: '#2196f3',
position: 'relative', position: 'relative',
@ -92,8 +92,6 @@ const Pulse: React.FC<PulseProps> = ({ timestamp, sx }) => {
animation: 'ripple-expand 1s ease-out 0.3s forwards', animation: 'ripple-expand 1s ease-out 0.3s forwards',
}; };
return ( return (
<> <>
<style> <style>

View File

@ -1,5 +1,6 @@
import { BackstoryMessage } from '../../Components/Message'; import { BackstoryMessage } from '../../Components/Message';
import { Query } from '../../Components/ChatQuery'; import { Query } from '../../Components/ChatQuery';
import { jsonrepair } from 'jsonrepair';
type StreamQueryOptions = { type StreamQueryOptions = {
query: Query; query: Query;
@ -66,7 +67,7 @@ const streamQueryResponse = (options: StreamQueryOptions) => {
let streaming_response = ''; let streaming_response = '';
const processLine = async (line: string) => { const processLine = async (line: string) => {
const update = JSON.parse(line); const update = JSON.parse(jsonrepair(line));
switch (update.status) { switch (update.status) {
case "streaming": case "streaming":

View File

@ -64,7 +64,7 @@ const GenerateCandidate = (props: BackstoryElementProps) => {
controllerRef.current = streamQueryResponse({ controllerRef.current = streamQueryResponse({
query: { 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: { agent_options: {
username: userRef.current.username, username: userRef.current.username,
filename: "profile.png" filename: "profile.png"
@ -74,7 +74,7 @@ const GenerateCandidate = (props: BackstoryElementProps) => {
sessionId, sessionId,
connectionBase, connectionBase,
onComplete: (msg) => { onComplete: (msg) => {
console.log({ msg, state: stateRef.current, prompt: promptRef.current || '' }); console.log(msg);
switch (msg.status) { switch (msg.status) {
case "partial": case "partial":
case "done": case "done":
@ -98,6 +98,9 @@ const GenerateCandidate = (props: BackstoryElementProps) => {
break; break;
default: default:
const data = JSON.parse(msg.response || ''); const data = JSON.parse(msg.response || '');
if (msg.status !== "heartbeat") {
console.log(data);
}
if (data.timestamp) { if (data.timestamp) {
setTimestamp(data.timestamp); setTimestamp(data.timestamp);
} else { } else {
@ -156,6 +159,7 @@ const GenerateCandidate = (props: BackstoryElementProps) => {
if (msg.status === "done") { if (msg.status === "done") {
setProcessing(false); setProcessing(false);
setCanGenImage(true); setCanGenImage(true);
setStatus('');
controllerRef.current = null; controllerRef.current = null;
stateRef.current = 0; stateRef.current = 0;
setTimeout(() => { setTimeout(() => {
@ -263,17 +267,17 @@ const GenerateCandidate = (props: BackstoryElementProps) => {
</Box> </Box>
} }
<Box sx={{display: "flex", flexDirection: "column"}}> <Box sx={{display: "flex", flexDirection: "column"}}>
<Box sx={{ display: "flex", flexDirection: "row"}}> <Box sx={{ display: "flex", flexDirection: "row", position: "relative" }}>
<Avatar <Avatar
src={hasProfile ? `/api/u/${user.username}/profile/${sessionId}` : ''} src={hasProfile ? `/api/u/${userRef.current.username}/profile/${sessionId}` : ''}
alt={`${user.full_name}'s profile`} alt={`${userRef.current.full_name}'s profile`}
sx={{ sx={{
width: 80, width: 80,
height: 80, height: 80,
border: '2px solid #e0e0e0', 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"}> <Tooltip title={"Generate Profile Picture"}>
<span style={{ display: "flex", flexGrow: 1 }}> <span style={{ display: "flex", flexGrow: 1 }}>
<Button <Button

View File

@ -96,7 +96,9 @@ class ImageGenerator(Agent):
logger.info("Beginning image generation...", self.user) 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.") 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" 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( async for message in generate_image(
message=message, message=message,
request=request request=request
@ -104,6 +106,10 @@ class ImageGenerator(Agent):
if message.status != "done": if message.status != "done":
yield message yield message
logger.info("Image generation done...") logger.info("Image generation done...")
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", []) images = self.user.get("images", [])
if self.filename not in images: if self.filename not in images:
images.append(self.filename) images.append(self.filename)

View File

@ -72,14 +72,14 @@ def flux_worker(pipe: Any, params: ImageRequest, status_queue: queue.Queue, task
# Simulate your pipe call with progress updates # Simulate your pipe call with progress updates
def status_callback(pipeline, step, timestep, callback_kwargs): def status_callback(pipeline, step, timestep, callback_kwargs):
# Send progress updates # 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_queue.put({
"status": "running", "status": "running",
"message": f"Processing step {step}/{params.iterations}", "message": f"Processing step {step+1}/{params.iterations} ({progress}%)",
"progress": progress "progress": progress
}) })
return callback_kwargs
# Replace this block with your actual Flux pipe call: # Replace this block with your actual Flux pipe call:
image = pipe( image = pipe(
@ -91,10 +91,10 @@ def flux_worker(pipe: Any, params: ImageRequest, status_queue: queue.Queue, task
callback_on_step_end=status_callback, callback_on_step_end=status_callback,
).images[0] ).images[0]
# Simulate image generation completion
gen_time = time.time() - start_gen_time gen_time = time.time() - start_gen_time
per_step_time = gen_time / params.iterations if params.iterations > 0 else 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) image.save(params.filepath)
# Final completion status # Final completion status
@ -104,10 +104,12 @@ def flux_worker(pipe: Any, params: ImageRequest, status_queue: queue.Queue, task
"progress": 100, "progress": 100,
"generation_time": gen_time, "generation_time": gen_time,
"per_step_time": per_step_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: except Exception as e:
logger.error(traceback.format_exc())
logger.error(e)
status_queue.put({ status_queue.put({
"status": "error", "status": "error",
"message": f"Generation failed: {str(e)}", "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)}', 'message': f'Server error: {str(e)}',
'error': str(e) 'error': str(e)
} }
logger.error(error_status)
yield error_status yield error_status
finally: finally:
@ -253,6 +256,8 @@ async def generate_image(message: Message, request: ImageRequest) -> AsyncGenera
async for update in async_generate_image(pipe, request): async for update in async_generate_image(pipe, request):
message.status = update.get("status", "thinking") message.status = update.get("status", "thinking")
message.response = json.dumps(update) # Merge properties from async_generate_image over the message... message.response = json.dumps(update) # Merge properties from async_generate_image over the message...
if message.status != "heartbeat":
logger.info(message.response)
yield message yield message
# Final result # Final result